Port gcd to structured (#57624)

Freey0 · facebook-github-bot · commit 7e51ac5ea72c · 2021-05-07T12:39:51.000-07:00
Summary: Pull Request resolved: #57624 Test Plan: Imported from OSS Reviewed By: VitalyFedyunin Differential Revision: D28224832 Pulled By: ezyang fbshipit-source-id: 30a8eba025c67d990103e49c03a396810f9d4006
diff --git a/aten/src/ATen/native/BinaryOps.cpp b/aten/src/ATen/native/BinaryOps.cpp
@@ -67,6 +67,14 @@ TORCH_META_FUNC(atan2) (const Tensor& self, const Tensor& other) {
   build_binary_float_op(maybe_get_output(), self, other);
 }
 
+// These are normal binary ops that preserve dtype
+#define CREATE_BINARY_META_FUNC(func)                                 \
+  TORCH_META_FUNC(func) (const Tensor& self, const Tensor& other) {   \
+    build_binary_op(maybe_get_output(), self, other);                 \
+  }
+
+CREATE_BINARY_META_FUNC(gcd);
+
 } // namespace meta
 
 
@@ -198,6 +206,13 @@ TORCH_IMPL_FUNC(special_xlog1py_out) (const Tensor& self, const Tensor& other, c
   xlog1py_stub(device_type(), *this);
 }
 
+#define CREATE_BINARY_TORCH_IMPL_FUNC(func)                                                    \
+TORCH_IMPL_FUNC(func##_out) (const Tensor& self, const Tensor& other, const Tensor& result) {  \
+  func##_stub(device_type(), *this);                                                           \
+}
+
+CREATE_BINARY_TORCH_IMPL_FUNC(gcd);
+
 Tensor special_xlog1py(const Scalar& x, const Tensor& y) {
   return at::special_xlog1py(wrapped_scalar_tensor(x), y);
 }
@@ -1062,21 +1077,6 @@ Tensor logaddexp2(const Tensor& self, const Tensor& other) {
   return at::logaddexp2_out(result, self, other);
 }
 
-Tensor& gcd_out(const Tensor& self, const Tensor& other, Tensor& result) {
-  auto iter = TensorIterator::binary_op(result, self, other);
-  gcd_stub(iter.device_type(), iter);
-  return result;
-}
-
-Tensor gcd(const Tensor& self, const Tensor& other) {
-  Tensor result = at::empty({0}, self.options());
-  return at::gcd_out(result, self, other);
-}
-
-Tensor& gcd_(Tensor& self, const Tensor& other) {
-  return at::gcd_out(self, self, other);
-}
-
 Tensor& lcm_out(const Tensor& self, const Tensor& other, Tensor& result) {
   auto iter = TensorIterator::binary_op(result, self, other);
   lcm_stub(iter.device_type(), iter);
diff --git a/aten/src/ATen/native/BinaryOps.h b/aten/src/ATen/native/BinaryOps.h
@@ -83,7 +83,7 @@ DECLARE_DISPATCH(binary_fn, mse_stub);
 DECLARE_DISPATCH(binary_fn, fmod_stub);
 DECLARE_DISPATCH(binary_fn, logaddexp_stub);
 DECLARE_DISPATCH(binary_fn, logaddexp2_stub);
-DECLARE_DISPATCH(binary_fn, gcd_stub);
+DECLARE_DISPATCH(structured_binary_fn, gcd_stub);
 DECLARE_DISPATCH(binary_fn, lcm_stub);
 DECLARE_DISPATCH(binary_fn, hypot_stub);
 DECLARE_DISPATCH(binary_fn, igamma_stub);
diff --git a/aten/src/ATen/native/cpu/BinaryOpsKernel.cpp b/aten/src/ATen/native/cpu/BinaryOpsKernel.cpp
@@ -858,7 +858,7 @@ void logaddexp2_kernel(TensorIterator& iter) {
   });
 }
 
-void gcd_kernel(TensorIterator& iter) {
+void gcd_kernel(TensorIteratorBase& iter) {
   AT_DISPATCH_INTEGRAL_TYPES(iter.dtype(), "gcd_cpu", [&]() {
       cpu_kernel(
           iter,
diff --git a/aten/src/ATen/native/cuda/GcdLcmKernel.cu b/aten/src/ATen/native/cuda/GcdLcmKernel.cu
@@ -10,7 +10,7 @@
 
 namespace at { namespace native {
 
-void gcd_kernel_cuda(TensorIterator& iter) {
+void gcd_kernel_cuda(TensorIteratorBase& iter) {
   AT_DISPATCH_INTEGRAL_TYPES(iter.dtype(), "gcd_cuda", [&]() {
     gpu_kernel(iter, [] GPU_LAMBDA (scalar_t a, scalar_t b) -> scalar_t {
       return calc_gcd(a, b);
diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml
@@ -1977,13 +1977,17 @@
     CPU: from_file
 
 - func: gcd.out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!)
+  structured: True
+  structured_inherits: TensorIteratorBase
   dispatch:
     CPU, CUDA: gcd_out
 
 - func: gcd(Tensor self, Tensor other) -> Tensor
+  structured_delegate: gcd.out
   variants: function, method
 
 - func: gcd_(Tensor(a!) self, Tensor other) -> Tensor(a!)
+  structured_delegate: gcd.out
   variants: function, method
 
 - func: lcm.out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!)

Original file line number	Diff line number	Diff line change
`@@ -858,7 +858,7 @@ void logaddexp2_kernel(TensorIterator& iter) {`
`858`	`858`	`});`
`859`	`859`	`}`
`860`	`860`
`861`		`-void gcd_kernel(TensorIterator& iter) {`
	`861`	`+void gcd_kernel(TensorIteratorBase& iter) {`
`862`	`862`	`AT_DISPATCH_INTEGRAL_TYPES(iter.dtype(), "gcd_cpu", [&]() {`
`863`	`863`	`cpu_kernel(`
`864`	`864`	`iter,`