Rewrite lerp operator to use TensorIterator and support compile-time vectorization. (#22038)

VitalyFedyunin · facebook-github-bot · commit fe580e850e68 · 2019-06-21T11:39:27.000-07:00
Summary: Get benefit from the compile time vectorization and multi-threading. Before: ```python In [1]: import torch In [2]: x = torch.randn(1000000) In [3]: y = torch.randn(1000000) In [4]: w = 0.7 In [5]: timeit torch.lerp(x, y, w) 2.29 ms ± 23.9 µs per loop (mean ± std. dev. of 7 runs, 100 loops each) ``` After: ```python In [1]: import torch In [2]: x = torch.randn(1000000) In [3]: y = torch.randn(1000000) In [4]: w = 0.7 In [5]: timeit torch.lerp(x, y, w) 452 µs ± 1.81 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each) ``` After with multi-processing: ```python In [1]: import torch In [2]: x = torch.randn(1000000) In [3]: y = torch.randn(1000000) In [4]: w = 0.7 In [5]: timeit torch.lerp(x, y, w) 167 µs ± 48.8 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each) ``` Pull Request resolved: #22038 Differential Revision: D15941468 Pulled By: VitalyFedyunin fbshipit-source-id: fa8a5126187df4e6c849452e035b00b22be25739
diff --git a/aten/src/ATen/native/Lerp.cpp b/aten/src/ATen/native/Lerp.cpp
@@ -1,111 +1,56 @@
+#include <ATen/native/Lerp.h>
+
 #include <ATen/ATen.h>
 #include <ATen/CPUApplyUtils.h>
 #include <ATen/NativeFunctions.h>
 #include <ATen/Dispatch.h>
 #include <ATen/ExpandUtils.h>
 
-namespace {
-template <typename scalar_t>
-void lerp_cpu(at::Tensor& ret, const at::Tensor& self, const at::Tensor& end, const at::Tensor& weight) {
-  at::CPU_tensor_apply4<scalar_t, scalar_t, scalar_t, scalar_t>(
-      ret, self, end, weight,
-      [](scalar_t& ret_val,
-         const scalar_t& self_val,
-         const scalar_t& end_val,
-         const scalar_t& weight_val) {
-        ret_val = (weight_val < 0.5) ?
-            self_val + weight_val * (end_val - self_val) : end_val - (end_val - self_val) * (1 - weight_val);
-      });
-}
-
-template <typename scalar_t>
-void lerp_cpu(at::Tensor& ret, const at::Tensor& self, const at::Tensor& end, scalar_t weight_val) {
-  at::CPU_tensor_apply3<scalar_t, scalar_t, scalar_t>(
-      ret, self, end,
-      [=](scalar_t& ret_val,
-         const scalar_t& self_val,
-         const scalar_t& end_val) {
-        ret_val = (weight_val < 0.5) ?
-            self_val + weight_val * (end_val - self_val) : end_val - (end_val - self_val) * (1 - weight_val);
-      });
-}
-
-} // namespace
-
 namespace at {
 namespace native {
 
 Tensor& lerp_cpu_tensor_out(Tensor& result, const Tensor& self,
                             const Tensor& end, const Tensor& weight) {
-  Tensor b_self, b_end, b_weight;
   TORCH_CHECK(weight.dim() <= std::max(self.dim(), end.dim()),
            "weight should be of dimension max(self.dim(), end.dim()) or lesser");
-  std::tie(b_self, b_end, b_weight) = expand_outplace(self, end, weight, "lerp_out_cpu");
-  result.resize_as_(b_self);
-  AT_DISPATCH_FLOATING_TYPES(self.scalar_type(), "lerp_out_cpu", [&]{
-    lerp_cpu<scalar_t>(result, b_self, b_end, b_weight);
-  });
+  lerp_kernel_tensor_weight(kCPU, result, self, end, weight);
   return result;
 }
 
 Tensor& lerp_cpu_scalar_out(Tensor& result, const Tensor& self,
                             const Tensor& end, Scalar weight) {
-  Tensor b_self, b_end;
-  std::tie(b_self, b_end) = expand_outplace(self, end, "lerp_out_cpu");
-  result.resize_as_(b_self);
-  AT_DISPATCH_FLOATING_TYPES(self.scalar_type(), "lerp_out_cpu", [&]{
-    lerp_cpu<scalar_t>(result, b_self, b_end, weight.to<scalar_t>());
-  });
+  lerp_kernel_scalar_weight(kCPU, result, self, end, weight);
   return result;
 }
 
 Tensor& lerp_cpu_tensor_(Tensor& self, const Tensor& end, const Tensor& weight) {
-  Tensor b_self, b_end, b_weight;
-  std::tie(b_self, b_end, b_weight) = expand_outplace(self, end, weight, "lerp__cpu");
-  TORCH_CHECK(b_self.sizes() == self.sizes(),
-           "output with shape ", self.sizes(),
-           " doesn't match the broadcast shape ", b_self.sizes());
   TORCH_CHECK(weight.dim() <= std::max(self.dim(), end.dim()),
            "weight should be of dimension max(self.dim(), end.dim()) or lesser");
-  AT_DISPATCH_FLOATING_TYPES(self.scalar_type(), "lerp__cpu", [&]{
-    lerp_cpu<scalar_t>(self, b_self, b_end, b_weight);
-  });
+  lerp_kernel_tensor_weight(kCPU, self, self, end, weight);
   return self;
 }
 
 Tensor& lerp_cpu_scalar_(Tensor& self, const Tensor& end, Scalar weight) {
-  Tensor b_self, b_end;
-  std::tie(b_self, b_end) = expand_outplace(self, end, "lerp__cpu");
-  TORCH_CHECK(b_self.sizes() == self.sizes(),
-           "output with shape ", self.sizes(),
-           " doesn't match the broadcast shape ", b_self.sizes());
-  AT_DISPATCH_FLOATING_TYPES(self.scalar_type(), "lerp__cpu", [&]{
-    lerp_cpu<scalar_t>(self, b_self, b_end, weight.to<scalar_t>());
-  });
+  lerp_kernel_scalar_weight(kCPU, self, self, end, weight);
   return self;
 }
 
 Tensor lerp_cpu_tensor(const Tensor& self, const Tensor& end, const Tensor& weight) {
-  Tensor b_self, b_end, b_weight;
   TORCH_CHECK(weight.dim() <= std::max(self.dim(), end.dim()),
            "weight should be of dimension max(self.dim(), end.dim()) or lesser");
-  std::tie(b_self, b_end, b_weight) = expand_outplace(self, end, weight, "lerp_cpu");
-  Tensor result = at::empty_like(b_self);
-  AT_DISPATCH_FLOATING_TYPES(result.scalar_type(), "lerp_cpu", [&]{
-    lerp_cpu<scalar_t>(result, b_self, b_end, b_weight);
-  });
+  Tensor result = at::empty({0}, self.options());
+  lerp_kernel_tensor_weight(kCPU, result, self, end, weight);
   return result;
 }
 
 Tensor lerp_cpu_scalar(const Tensor& self, const Tensor& end, Scalar weight) {
-  Tensor b_self, b_end;
-  std::tie(b_self, b_end) = expand_outplace(self, end, "lerp_cpu");
-  Tensor result = at::empty_like(b_self);
-  AT_DISPATCH_FLOATING_TYPES(result.scalar_type(), "lerp_cpu", [&]{
-    lerp_cpu<scalar_t>(result, b_self, b_end, weight.to<scalar_t>());
-  });
+  Tensor result = at::empty({0}, self.options());
+  lerp_kernel_scalar_weight(kCPU, result, self, end, weight);
   return result;
 }
 
+DEFINE_DISPATCH(lerp_kernel_scalar_weight);
+DEFINE_DISPATCH(lerp_kernel_tensor_weight);
+
 } // namespace native
 } // namespace at
diff --git a/aten/src/ATen/native/Lerp.h b/aten/src/ATen/native/Lerp.h
@@ -0,0 +1,25 @@
+#pragma once
+
+#include <ATen/ATen.h>
+#include <ATen/native/DispatchStub.h>
+
+namespace at {
+namespace native {
+
+using lerp_fn_scalar = void (*)(
+    at::Tensor& ret,
+    const at::Tensor& self,
+    const at::Tensor& end,
+    Scalar weight);
+
+using lerp_fn_tensor = void (*)(
+    at::Tensor& ret,
+    const at::Tensor& self,
+    const at::Tensor& end,
+    const at::Tensor& weights);
+
+DECLARE_DISPATCH(lerp_fn_scalar, lerp_kernel_scalar_weight);
+DECLARE_DISPATCH(lerp_fn_tensor, lerp_kernel_tensor_weight);
+
+} // namespace native
+} // namespace at
diff --git a/aten/src/ATen/native/cpu/LerpKernel.cpp b/aten/src/ATen/native/cpu/LerpKernel.cpp
@@ -0,0 +1,62 @@
+#include <ATen/ATen.h>
+
+#include <ATen/Dispatch.h>
+#include <ATen/native/Lerp.h>
+#include <ATen/native/TensorIterator.h>
+#include <ATen/native/cpu/Loops.h>
+
+namespace at {
+namespace native {
+namespace {
+
+static void lerp_kernel_scalar(
+    Tensor& ret,
+    const Tensor& self,
+    const Tensor& end,
+    Scalar weight) {
+  auto builder = at::TensorIterator::Builder();
+  builder.add_output(ret);
+  builder.add_input(self);
+  builder.add_input(end);
+  auto iter = builder.build();
+  AT_DISPATCH_FLOATING_TYPES(ret.scalar_type(), "lerp_kernel_scalar", [&] {
+    scalar_t weight_val = weight.to<scalar_t>();
+    at::native::cpu_kernel(
+        *iter,
+        [weight_val](scalar_t self_val, scalar_t end_val) {
+          return (weight_val < 0.5)
+              ? self_val + weight_val * (end_val - self_val)
+              : end_val - (end_val - self_val) * (1 - weight_val);
+        });
+  });
+}
+
+static void lerp_kernel_tensor(
+    Tensor& ret,
+    const Tensor& self,
+    const Tensor& end,
+    const Tensor& weights) {
+  auto builder = at::TensorIterator::Builder();
+  builder.add_output(ret);
+  builder.add_input(self);
+  builder.add_input(end);
+  builder.add_input(weights);
+  auto iter = builder.build();
+  AT_DISPATCH_FLOATING_TYPES(ret.scalar_type(), "lerp_kernel_tensor", [&] {
+    at::native::cpu_kernel(
+        *iter,
+        [](scalar_t self_val, scalar_t end_val, scalar_t weight_val) {
+          return (weight_val < 0.5)
+              ? self_val + weight_val * (end_val - self_val)
+              : end_val - (end_val - self_val) * (1 - weight_val);
+        });
+  });
+}
+
+} // anonymous namespace
+
+REGISTER_DISPATCH(lerp_kernel_scalar_weight, &lerp_kernel_scalar);
+REGISTER_DISPATCH(lerp_kernel_tensor_weight, &lerp_kernel_tensor);
+
+} // namespace native
+} // namespace at
