Foreach Binary Test Refactor (#59907)

crcrpar · facebook-github-bot · commit fac744e1161a · 2021-07-06T11:49:38.000-07:00
Summary: Related: #58833 ## Changes I'm a bit concerned - binary ops with one tensorlist and one scalarlist support complex dtypes. To realize this, I added a specialization of [`TensorListScalarListMetadata<c10::complex<double>, 1>` ](https://github.com/pytorch/pytorch/pull/59907/files#diff-131eb9b310905b15b3528da6a23e542a3a3aa952bc88f7423c98a23a8a28cca1R49). This might be out of the scope of this pull request. cc ptrblck ngimel mcarilli Pull Request resolved: #59907 Reviewed By: mruberry Differential Revision: D29551001 Pulled By: ngimel fbshipit-source-id: 46b25fdba85dd4d6332a77b27376fe96cd422384
diff --git a/aten/src/ATen/native/BinaryOps.h b/aten/src/ATen/native/BinaryOps.h
@@ -29,7 +29,7 @@ inline void sub_check(const Tensor& self, const Tensor& other) {
 
 inline void sub_check(const Tensor& self, const Scalar& scalar) {
   TORCH_CHECK(self.scalar_type() != kBool || !scalar.isBoolean(),
-              "Subtraction, the `-` operator, with two bool tensors is not supported."
+              "Subtraction, the `-` operator, with two bool tensors is not supported. "
               "Use the `^` or `logical_xor()` operator instead.")
   TORCH_CHECK(self.scalar_type() != kBool && !scalar.isBoolean(),
               "Subtraction, the `-` operator, with a bool tensor is not supported. "
diff --git a/aten/src/ATen/native/ForeachUtils.h b/aten/src/ATen/native/ForeachUtils.h
@@ -65,20 +65,6 @@ void check_foreach_api_restrictions(TensorList tensors1, TensorList tensors2, Te
 // - All tensors must be non-overlapping and dense
 // - Resulting tensor must have the same dtype as the input one
 
-// TODO(mkozuki): Consider whether we really need this function or not.
-// Note that, there is a possibility that foreach fastpath supports type promotion in the future,
-// which might complicate the functionality this function should provides.
-// However, as of now, the check of division op with integer inputs is duplicated.
-// `check_fast_path_restrictions` does the same thing in it before calling this function.
-bool will_promote_tensor(const Tensor& tensor, const Scalar& scalar, bool does_op_promote_integer_inputs_to_float = false) {
-  // In case of division, integer inputs will result in float
-  if (does_op_promote_integer_inputs_to_float &&
-      at::isIntegralType(tensor.scalar_type(), /* includeBool */ true)) {
-    return true;
-  }
-  return tensor.scalar_type() != at::native::result_type(scalar, tensor);
-}
-
 // Please, make sure to call check_foreach_api_restrictions before calling this method.
 // There is a set of preconditions that have to be satisfied.
 bool check_fast_path_restrictions(
@@ -114,29 +100,21 @@ bool check_fast_path_restrictions(
 
     // This function has already checked that `tensorList[j][i]` for all j, i has the same dtype
     // using `is_tensor_okay` function above.
-    // checked by `check_foreach_api_restrictions`).
     // This means we only need to check if {tensorList[0][0], tensorList[0][1], tensorList[0][2], ...}
     // do type promotion with scalarLIst.
     for (const auto i : c10::irange(tensorLists[0].size())) {
+      // For division, integer inputs will result in float.
       if (does_op_promote_integer_inputs_to_float) {
         if (at::isIntegralType(tensorLists[0][i].scalar_type(), /*includeBool*/ true)) {
           return false;
         }
       }
-
-      if (scalarList.size() == 1) {
-        if (will_promote_tensor(tensorLists[0][i], scalarList[0])) {
-          return false;
-        }
-      } else if (scalarList.size() > 1) {
-        // FIXME(mkozuki): Consider specializing `TensorListScalarListMetadata` for complex dtypes
-        // to access the following comment.
-        // Complex scalar list is not supported due to the limit for kernel launch argument (4KB)
-        if (scalarList[i].isComplex()) {
-          return false;
-        }
-
-        if (will_promote_tensor(tensorLists[0][i], scalarList[i])) {
+      if (scalarList.size() > 0) {
+        const auto& scalar = scalarList.size() == 1 ? scalarList[0] : scalarList[i];
+        const auto& tensor = tensorLists[0][i];
+        // note(mkozuki): This check might be responsible for `_foreach_add(bool_tensors, bool_tensors)`
+        // being pushed to slow path.
+        if (tensor.scalar_type() != at::native::result_type(scalar, tensor)) {
           return false;
         }
       }
diff --git a/aten/src/ATen/native/cuda/ForeachBinaryOpScalarList.cu b/aten/src/ATen/native/cuda/ForeachBinaryOpScalarList.cu
@@ -17,7 +17,7 @@ std::vector<Tensor> foreach_binary_op(TensorList tensors, at::ArrayRef<Scalar> s
     tensor_lists.emplace_back(tensors.vec());
     tensor_lists.emplace_back(vec_res);
 
-    AT_DISPATCH_ALL_TYPES_AND3(kBFloat16, kHalf, kBool, tensors[0].scalar_type(), "foreach_binary_op_scalarlist_cuda", [&]() {
+    AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND3(kBFloat16, kHalf, kBool, tensors[0].scalar_type(), "foreach_binary_op_scalarlist_cuda", [&]() {
         using opmath_t = get_opmath_t<scalar_t>::opmath_t;
         multi_tensor_apply<2, opmath_t>(tensor_lists,
                                         scalars,
@@ -36,7 +36,7 @@ void foreach_binary_op_(TensorList tensors, at::ArrayRef<Scalar> scalars) {
     std::vector<std::vector<at::Tensor>> tensor_lists;
     tensor_lists.emplace_back(tensors.vec());
 
-    AT_DISPATCH_ALL_TYPES_AND3(kBFloat16, kHalf, kBool, tensors[0].scalar_type(), "foreach_binary_op_scalarlist_cuda_", [&]() {
+    AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND3(kBFloat16, kHalf, kBool, tensors[0].scalar_type(), "foreach_binary_op_scalarlist_cuda_", [&]() {
         using opmath_t = get_opmath_t<scalar_t>::opmath_t;
         multi_tensor_apply<1, opmath_t>(tensor_lists,
                                         scalars,
diff --git a/aten/src/ATen/native/cuda/MultiTensorApply.cuh b/aten/src/ATen/native/cuda/MultiTensorApply.cuh
@@ -46,6 +46,18 @@ template<typename scalar_vals_t, int n> struct TensorListScalarListMetadata
   int block_to_chunk[depth_to_max_blocks[n-1]];
 };
 
+// note(mkozuki): `n` of 96 and `scalar_vals_t` of `c10::complex<double>`
+// violates the cuda kernel argument size limitation of 4kb.
+// 80 is a number that does not violate this limitation.
+template<> struct TensorListScalarListMetadata<c10::complex<double>, 1>
+{
+  void* addresses[1][80];
+  int numel_for_tensor[80];
+  c10::complex<double> scalar_vals[80];
+  unsigned char block_to_tensor[depth_to_max_blocks[1-1]];
+  int block_to_chunk[depth_to_max_blocks[1-1]];
+};
+
 template<typename T, typename U, typename... ArgTypes>
 C10_LAUNCH_BOUNDS_1(kBlockSize)
 __global__ void
diff --git a/test/test_foreach.py b/test/test_foreach.py
diff --git a/torch/testing/_internal/common_methods_invocations.py b/torch/testing/_internal/common_methods_invocations.py
