[Release/2.2] Fix NULL dereference in binary CPU ops (#115470)

malfet · web-flow · commit e57f089704ae · 2023-12-08T19:33:13.000-08:00
Targeted fix for #113037 A more fundamental one, where those functions are not even called for empty tensors are coming later Cherry-pick of release #115183 into release/2.2 branch (cherry picked from commit b56b002)
diff --git a/aten/src/ATen/native/cpu/BinaryOpsKernel.cpp b/aten/src/ATen/native/cpu/BinaryOpsKernel.cpp
@@ -130,7 +130,7 @@ void mul_kernel(TensorIteratorBase& iter) {
           using comp_t = c10::complex<float>;
           return comp_t{a} * comp_t{b};
         });
-  } else if (iter.is_scalar(2) && at::isReducedFloatingType(dtype)) {
+  } else if (iter.is_scalar(2) && iter.data_ptr(2) != nullptr && at::isReducedFloatingType(dtype)) {
     AT_DISPATCH_REDUCED_FLOATING_TYPES(dtype, "mul_cpu_reduced_float", [&]() {
       using opmath_t = at::opmath_type<scalar_t>;
       opmath_t b = iter.original_scalar_value<opmath_t>(2);
@@ -162,7 +162,7 @@ void mul_kernel(TensorIteratorBase& iter) {
 
 void div_true_kernel(TensorIteratorBase& iter) {
   const auto dtype = iter.common_dtype();
-  if (iter.is_scalar(2) && at::isReducedFloatingType(dtype)) {
+  if (iter.is_scalar(2) && iter.data_ptr(2) != nullptr && at::isReducedFloatingType(dtype)) {
     AT_DISPATCH_REDUCED_FLOATING_TYPES(dtype, "div_cpu_reduced_float", [&]() {
       using opmath_t = at::opmath_type<scalar_t>;
       opmath_t b = iter.original_scalar_value<opmath_t>(2);
@@ -208,7 +208,7 @@ void div_trunc_kernel(TensorIteratorBase& iter) {
         return a / b;
       });
     });
-  } else if (iter.is_scalar(2) && at::isReducedFloatingType(dtype)) {
+  } else if (iter.is_scalar(2) && iter.data_ptr(2) != nullptr && at::isReducedFloatingType(dtype)) {
     AT_DISPATCH_REDUCED_FLOATING_TYPES(
         dtype, "div_trunc_cpu_reduced_float", [&]() {
           using opmath_t = at::opmath_type<scalar_t>;
@@ -283,7 +283,7 @@ void div_floor_kernel(TensorIteratorBase& iter) {
     });
   } else {
     // See NOTE: [Floor Division in Python]
-    if (iter.is_scalar(2) && at::isReducedFloatingType(dtype)) {
+    if (iter.is_scalar(2) && iter.data_ptr(2) != nullptr && at::isReducedFloatingType(dtype)) {
       AT_DISPATCH_REDUCED_FLOATING_TYPES(
           dtype, "div_floor_cpu_reduced_float", [&]() {
             using opmath_t = at::opmath_type<scalar_t>;
diff --git a/test/test_foreach.py b/test/test_foreach.py
@@ -349,13 +349,17 @@ def _pointwise_test(
     @dtypes(*all_types_and_complex_and(torch.half, torch.bfloat16))
     def test_add_scalar_with_empty_list_and_empty_tensor(self, device, dtype):
         # TODO: enable empty list case
-        for tensors in [[torch.randn([0], device=device, dtype=dtype)]]:
+        for tensors in [[torch.randn([0], device=device, dtype=dtype)],
+                        [torch.empty_strided((0, 1), (0, 0), dtype=dtype, device=device)]]:
             res = torch._foreach_add(tensors, 1)
             self.assertEqual(res, tensors)
 
             torch._foreach_add_(tensors, 1)
             self.assertEqual(res, tensors)
 
+            # Regression test for https://github.com/pytorch/pytorch/issues/113156
+            torch._foreach_mul_(tensors, 1)
+
     @ops(
         filter(lambda op: op.supports_out, foreach_binary_op_db),
         dtypes=OpDTypes.supported,
diff --git a/test/test_numpy_interop.py b/test/test_numpy_interop.py
@@ -482,6 +482,21 @@ def test_numpy_scalar_cmp(self, device, dtype):
                 else:
                     self.assertTrue(t == a)
 
+    @onlyCPU
+    def test_empty_tensors_interop(self, device):
+        x = torch.rand((), dtype=torch.float16)
+        y = torch.tensor(np.random.rand(0), dtype=torch.float16)
+        # Same can be achieved by running
+        # y = torch.empty_strided((0,), (0,), dtype=torch.float16)
+
+        # Regression test for https://github.com/pytorch/pytorch/issues/115068
+        self.assertEqual(torch.true_divide(x, y).shape, y.shape)
+        # Regression test for https://github.com/pytorch/pytorch/issues/115066
+        self.assertEqual(torch.mul(x, y).shape, y.shape)
+        # Regression test for https://github.com/pytorch/pytorch/issues/113037
+        self.assertEqual(torch.div(x, y, rounding_mode='floor').shape, y.shape)
+
+
 instantiate_device_type_tests(TestNumPyInterop, globals())
 
 if __name__ == '__main__':
