[MPS] Add fixes for div with floor (#95869)

kulinseth · web-flow · commit c4572aa1b72d · 2023-03-02T12:36:02.000-08:00
* [MPS] Add fixes for div with floor and raise error for div_trunc (#95769) Fixes #ISSUE_NUMBER Pull Request resolved: #95769 Approved by: https://github.com/DenisVieriu97 * Add back the unittest skip for MacOS 12.
diff --git a/aten/src/ATen/native/mps/MPSGraphVenturaOps.h b/aten/src/ATen/native/mps/MPSGraphVenturaOps.h
@@ -138,4 +138,7 @@ typedef NS_ENUM(NSUInteger, MPSGraphResizeNearestRoundingMode)
                                      nearestRoundingMode:(MPSGraphResizeNearestRoundingMode) nearestRoundingMode
                                            constantValue:(double) constantValue
                                                     name:(NSString * _Nullable) name;
+- (MPSGraphTensor * _Nonnull) truncateWithTensor:(MPSGraphTensor * _Nonnull) tensor
+                                            name:(NSString * _Nullable) name;
+
 @end
diff --git a/aten/src/ATen/native/mps/operations/BinaryOps.mm b/aten/src/ATen/native/mps/operations/BinaryOps.mm
@@ -177,10 +177,6 @@ void div_mode_template(const Tensor& self, const Tensor& other,
                        c10::optional<c10::string_view> rounding_mode,
                        const Tensor& output, const string op_name)
 {
-  if(rounding_mode.has_value() && *rounding_mode == "floor"){
-    TORCH_CHECK(self.scalar_type() != ScalarType::Long,
-                "MPS: does not support floor_divide op with int64 input");
-  }
   BinaryOpBlock div_mode_op_block = ^BinaryOpFn(cachedGraph, primaryCastTensor, secondaryCastTensor) {
     MPSGraph* mpsGraph = cachedGraph->graph();
     bool isFloatInput = ([primaryCastTensor dataType] & MPSDataTypeFloatBit) != 0;
diff --git a/aten/src/ATen/native/mps/operations/UnaryOps.mm b/aten/src/ATen/native/mps/operations/UnaryOps.mm
@@ -75,15 +75,20 @@ void unary_op(const Tensor& self, const Tensor& output, std::string op_name, Una
     return inputTensor;
   }
 
-  MPSGraphTensor* zeroTensor = [mpsGraph constantWithScalar:0.0
-                                                   dataType:inputTensor.dataType];
-  MPSGraphTensor* predicateTensor = [mpsGraph lessThanWithPrimaryTensor:inputTensor
-                                                        secondaryTensor:zeroTensor
-                                                                    name:nil];
-  return [mpsGraph selectWithPredicateTensor:predicateTensor
-                         truePredicateTensor:[mpsGraph ceilWithTensor :inputTensor name:nil]
-                        falsePredicateTensor:[mpsGraph floorWithTensor:inputTensor name:nil]
-                                        name:nil];
+  if(!is_macos_13_or_newer()) {
+    MPSGraphTensor* zeroTensor = [mpsGraph constantWithScalar:0.0
+                                                    dataType:inputTensor.dataType];
+    MPSGraphTensor* predicateTensor = [mpsGraph lessThanWithPrimaryTensor:inputTensor
+                                                          secondaryTensor:zeroTensor
+                                                                      name:nil];
+    return [mpsGraph selectWithPredicateTensor:predicateTensor
+                          truePredicateTensor:[mpsGraph ceilWithTensor :inputTensor name:nil]
+                          falsePredicateTensor:[mpsGraph floorWithTensor:inputTensor name:nil]
+                                          name:nil];
+  } else {
+    return [mpsGraph truncateWithTensor:inputTensor
+                                   name:nil];
+  }
 };
 
 } // namespace mps
diff --git a/test/test_mps.py b/test/test_mps.py
@@ -2344,10 +2344,9 @@ def test_full_bugs(self):
     # See https://github.com/pytorch/pytorch/issues/84995
     def test_div_bugs(self):
         for (dtype, mode) in itertools.product(integral_types(), ['trunc', 'floor']):
-            if dtype != torch.int64:
-                x = torch.tensor(list(range(1, 11)), device='mps', dtype=dtype)
-                y = torch.div(x, 101, rounding_mode=mode)
-                self.assertEqual(y.sum(), 0)
+            x = torch.tensor(list(range(1, 11)), device='mps', dtype=dtype)
+            y = torch.div(x, 101, rounding_mode=mode)
+            self.assertEqual(y.sum(), 0)
 
     # See https://github.com/pytorch/pytorch/issues/82663
     def test_bool_expand(self):
