add logaddexp2

pytorch · cloudhan · May 13, 2020 · May 13, 2020 · May 14, 2020 · May 14, 2020
commit 8bbe55de733bbc5091372ed3c8dee0f033839de7
diff --git a/aten/src/ATen/native/BinaryOps.cpp b/aten/src/ATen/native/BinaryOps.cpp
@@ -37,6 +37,7 @@ DEFINE_DISPATCH(min_elementwise_stub);
 DEFINE_DISPATCH(fmod_stub);
 DEFINE_DISPATCH(fmod_scalar_stub);
 DEFINE_DISPATCH(logaddexp_stub);
+DEFINE_DISPATCH(logaddexp2_stub);
 
 Tensor& add_out(Tensor& result, const Tensor& self, const Tensor& other, Scalar alpha) {
   auto iter = TensorIterator::binary_op(result, self, other,
@@ -754,6 +755,17 @@ Tensor logaddexp(const Tensor& self, const Tensor& other) {
   return at::logaddexp_out(result, self, other);
 }
 
+Tensor& logaddexp2_out(Tensor& result, const Tensor& self, const Tensor& other) {
+  auto iter = TensorIterator::binary_op(result, self, other);
+  logaddexp2_stub(iter.device_type(), iter);
+  return result;
+}
+
+Tensor logaddexp2(const Tensor& self, const Tensor& other) {
+  Tensor result = at::empty({0}, self.options());
+  return at::logaddexp2_out(result, self, other);
+}
+
 Tensor true_divide(const Tensor& self, Scalar divisor) {
   return self.true_divide(wrapped_scalar_tensor(divisor)); // redispatch!
 }

diff --git a/aten/src/ATen/native/BinaryOps.h b/aten/src/ATen/native/BinaryOps.h
@@ -56,5 +56,6 @@ DECLARE_DISPATCH(binary_fn, mse_stub);
 DECLARE_DISPATCH(binary_fn, fmod_stub);
 DECLARE_DISPATCH(binary_fn_alpha, fmod_scalar_stub);
 DECLARE_DISPATCH(binary_fn, logaddexp_stub);
+DECLARE_DISPATCH(binary_fn, logaddexp2_stub);
 
 }} // namespace at::native
diff --git a/aten/src/ATen/native/cpu/BinaryOpsKernel.cpp b/aten/src/ATen/native/cpu/BinaryOpsKernel.cpp
@@ -630,6 +630,23 @@ void logaddexp_kernel(TensorIterator& iter) {
   });
 }
 
+void logaddexp2_kernel(TensorIterator& iter) {
+  AT_DISPATCH_FLOATING_TYPES(iter.dtype(), "logaddexp2_cpu", [&]() {
+    cpu_kernel_vec(
+        iter,
+        [=](scalar_t a, scalar_t b) -> scalar_t {
+          scalar_t m = std::max(a, b);
+          return m + std::log2((scalar_t)(1.0) + std::pow((scalar_t)(2), -std::abs(a - b)));
+        },
+        [=](Vec256<scalar_t> a, Vec256<scalar_t> b) {
+          Vec256<scalar_t> one(1.0);
+          Vec256<scalar_t> two(2.0);
+          Vec256<scalar_t> m = maximum(a, b);
+          return m + (one + two.pow((a - b).abs().neg())).log2();
+        });
+  });
+}
+
 } // anonymous namespace
 
 
@@ -662,5 +679,6 @@ REGISTER_DISPATCH(mse_stub, &mse_kernel);
 REGISTER_DISPATCH(fmod_stub, &fmod_kernel);
 REGISTER_DISPATCH(fmod_scalar_stub, &fmod_scalar_kernel);
 REGISTER_DISPATCH(logaddexp_stub, &logaddexp_kernel);
+REGISTER_DISPATCH(logaddexp2_stub, &logaddexp2_kernel);
 
 }} // namespace at::native
diff --git a/aten/src/ATen/native/cuda/BinaryMiscOpsKernels.cu b/aten/src/ATen/native/cuda/BinaryMiscOpsKernels.cu
@@ -48,9 +48,19 @@ void logaddexp_kernel_cuda(TensorIterator& iter) {
   });
 }
 
+void logaddexp2_kernel_cuda(TensorIterator& iter) {
+  AT_DISPATCH_FLOATING_TYPES_AND_HALF(iter.dtype(), "logaddexp2_cuda", [&]() {
+    gpu_kernel(iter, [] GPU_LAMBDA (scalar_t a, scalar_t b) -> scalar_t {
+      scalar_t m = ::max(a, b);
+      return m + ::log2((scalar_t)(1.0) + ::pow((scalar_t)(2.0), -::abs(a - b)));
+    });
+  });
+}
+
 REGISTER_DISPATCH(atan2_stub, &atan2_kernel_cuda);
 REGISTER_DISPATCH(smooth_l1_stub, &smooth_l1_kernel_cuda);
 REGISTER_DISPATCH(mse_stub, &mse_kernel_cuda);
 REGISTER_DISPATCH(logaddexp_stub, &logaddexp_kernel_cuda);
+REGISTER_DISPATCH(logaddexp2_stub, &logaddexp2_kernel_cuda);
 
 }} // namespace at::native
diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml
@@ -1770,6 +1770,14 @@
   supports_named_tensor: True
   variants: method, function
 
+- func: logaddexp2.out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!)
+  supports_named_tensor: True
+
+- func: logaddexp2(Tensor self, Tensor other) -> Tensor
+  use_c10_dispatcher: full
+  supports_named_tensor: True
+  variants: method, function
+
 - func: logdet(Tensor self) -> Tensor
   use_c10_dispatcher: full
   variants: function, method