[BE] Migrate all add/sub ops to Metal kernels (#152510)

malfet · pytorchmergebot · commit 52cbcac640dc · 2025-05-01T15:35:57.000Z
As typecasting harness shoudl take care of all permutations Fix bug in `exec_binary_kernel` where it was not properly downcasting CPU double/complexDouble scalars to floats Fixes #152582 Pull Request resolved: #152510 Approved by: https://github.com/Skylion007, https://github.com/dcci, https://github.com/cyyever ghstack dependencies: #152443, #152466, #152479, #152504, #152485
diff --git a/aten/src/ATen/native/mps/OperationUtils.mm b/aten/src/ATen/native/mps/OperationUtils.mm
@@ -1022,13 +1022,30 @@ static dispatch_data_t getSectionData(const std::string& name) {
 void MetalShaderLibrary::exec_binary_kernel(TensorIteratorBase& iter,
                                             const std::string& name,
                                             std::optional<c10::Scalar> alpha) {
-  TORCH_CHECK(iter.common_dtype() != at::kDouble, "float64 is not supported on MPS");
+  // TODO: Figure a better place to downcast double scalars (probably in tensor iterator itself?)
+  // Right now running something like 1.0-torch.rand(5, device='mps') will create iterator with
+  // double as common dtype (because Python floating point are always 64-bit values)
+  TORCH_CHECK(iter.output().scalar_type() != at::kDouble, "float64 is not supported on MPS");
   TORCH_CHECK(iter.can_use_32bit_indexing(), "Can't be indexed using 32-bit iterator");
 
+  auto convert_double_scalar = [](Tensor& t) {
+    if (t.dim() != 0) {
+      return;
+    }
+    if (t.scalar_type() == kDouble) {
+      t = t.to(kFloat);
+    } else if (t.scalar_type() == kComplexDouble) {
+      t = t.to(kComplexFloat);
+    }
+  };
+
   Tensor input = iter.input(0);
   Tensor other = iter.input(1);
   Tensor out = iter.output();
 
+  convert_double_scalar(input);
+  convert_double_scalar(other);
+
   id<MTLDevice> device = MPSDevice::getInstance()->device();
   MPSStream* mpsStream = getCurrentMPSStream();
   const uint32_t nDim = iter.ndim();
diff --git a/aten/src/ATen/native/mps/operations/BinaryKernel.mm b/aten/src/ATen/native/mps/operations/BinaryKernel.mm
@@ -64,8 +64,13 @@ void binary_op_kernel(const std::string func_name,
     return;
   }
 
-  auto iter =
-      TensorIteratorConfig().add_output(output).add_input(input).add_input(other).check_all_same_dtype(false).build();
+  auto iter = TensorIteratorConfig()
+                  .allow_cpu_scalars(true)
+                  .add_output(output)
+                  .add_input(input)
+                  .add_input(other)
+                  .check_all_same_dtype(false)
+                  .build();
 
   lib.exec_binary_kernel(iter, func_name, alpha);
 }
diff --git a/aten/src/ATen/native/mps/operations/BinaryOps.mm b/aten/src/ATen/native/mps/operations/BinaryOps.mm
@@ -45,7 +45,7 @@
 struct BinaryOpCachedGraph : public MPSCachedGraph {
   BinaryOpCachedGraph(MPSGraph* graph) : MPSCachedGraph(graph) {}
   MPSGraphTensor *primaryTensor = nil, *secondaryTensor = nil;
-  MPSGraphTensor *alphaTensor = nil, *outputTensor = nil;
+  MPSGraphTensor* outputTensor = nil;
 };
 
 typedef MPSGraphTensor* (^BinaryOpBlock)(BinaryOpCachedGraph*, MPSGraphTensor*, MPSGraphTensor*);
@@ -60,10 +60,8 @@ static inline Tensor legacy_complex_as_view(const Tensor& t) {
   return at::view_as_real(t.dim() != 0 ? t : t.to(kMPS));
 }
 
-// alpha is always 1.0 except when this function is called from add_sub_lerp_template()
 static void binaryOpTensor(const Tensor& self,
                            const Tensor& other,
-                           const Scalar& alpha,
                            const Tensor& output_,
                            std::string op_name,
                            BinaryOpBlock binaryBlock) {
@@ -147,7 +145,6 @@ static void binaryOpTensor(const Tensor& self,
     Placeholder otherPlaceholder;
     MPSScalar self_scalar;
     MPSScalar other_scalar;
-    MPSScalar alpha_scalar;
 
     if (is_self_scalar && !self.is_mps()) {
       self_scalar = getMPSScalar(self.item(), inputDataType);
@@ -172,12 +169,6 @@ static void binaryOpTensor(const Tensor& self,
       feeds[otherPlaceholder.getMPSGraphTensor()] = otherPlaceholder.getMPSGraphTensorData();
     }
 
-    // 'cachedGraph->alphaTensor' is not nil only if add_sub_lerp_template() was called with an alpha value != 1.0
-    if (cachedGraph->alphaTensor) {
-      alpha_scalar = getMPSScalar(alpha, common_dtype);
-      feeds[cachedGraph->alphaTensor] = getMPSGraphTensorFromScalar(mpsStream, alpha_scalar);
-    }
-
     Placeholder outputPlaceholder = Placeholder(cachedGraph->outputTensor, needsCopyToOutput ? output : output_);
     runMPSGraph(mpsStream, cachedGraph->graph(), feeds, outputPlaceholder);
 
@@ -189,11 +180,10 @@ static void binaryOpTensor(const Tensor& self,
 
 static void binaryOpScalar(const Tensor& self,
                            const Scalar& other,
-                           const Scalar& alpha,
                            const Tensor& output,
                            std::string op_name,
                            BinaryOpBlock binaryBlock) {
-  binaryOpTensor(self, wrapped_scalar_tensor(other), alpha, output, op_name, binaryBlock);
+  binaryOpTensor(self, wrapped_scalar_tensor(other), output, op_name, binaryBlock);
 }
 
 static void div_mode_template(const Tensor& self,
@@ -246,7 +236,6 @@ static void div_mode_template(const Tensor& self,
   };
   binaryOpTensor(self,
                  other,
-                 Scalar(1.0),
                  output,
                  op_name + "_mps:" + (rounding_mode.has_value() ? c10::str(*rounding_mode) : ""),
                  div_mode_op_block);
@@ -272,100 +261,51 @@ static void add_sub_lerp_template(const Tensor& self,
     return;
   }
 
-  auto self_complex = c10::isComplexType(self.scalar_type());
-  auto other_complex = c10::isComplexType(other.scalar_type());
-  auto commonDtype = at::result_type(self, other);
-  if (self.is_mps() && other.is_mps() && (output.scalar_type() == commonDtype) && (self_complex == other_complex)) {
-    if (alpha_has_value) {
-      at::native::alpha_check(commonDtype, alpha);
-      mps::binary_op_kernel(op_name + "_alpha", self, other, output, alpha);
-    } else {
-      mps::binary_op_kernel(op_name, self, other, output);
-    }
-    return;
+  if (alpha_has_value) {
+    auto commonDtype = at::result_type(self, other);
+    at::native::alpha_check(commonDtype, alpha);
+    mps::binary_op_kernel(op_name + "_alpha", self, other, output, alpha);
+  } else {
+    mps::binary_op_kernel(op_name, self, other, output);
   }
-
-  BinaryOpBlock add_sub_lerp_op_block = ^BinaryOpFn(cachedGraph, primaryCastTensor, secondaryCastTensor) {
-    MPSGraph* mpsGraph = cachedGraph->graph();
-    MPSGraphTensor* secondaryTensor = secondaryCastTensor;
-
-    if (op_name == "lerp") {
-      secondaryCastTensor = [mpsGraph subtractionWithPrimaryTensor:secondaryCastTensor
-                                                   secondaryTensor:primaryCastTensor
-                                                              name:nil];
-    }
-
-    // if alpha is 1.0, then we don't bother adding another multiply to graph
-    if (alpha_has_value) {
-      auto commonDtype = c10::promoteTypes(self.scalar_type(), other.scalar_type());
-      cachedGraph->alphaTensor = mpsGraphRankedPlaceHolder(mpsGraph, getMPSScalarType(commonDtype), @[ @1 ]);
-      secondaryTensor = [mpsGraph multiplicationWithPrimaryTensor:secondaryCastTensor
-                                                  secondaryTensor:cachedGraph->alphaTensor
-                                                             name:nil];
-    }
-    if (op_name == "add" || op_name == "lerp")
-      return [mpsGraph additionWithPrimaryTensor:primaryCastTensor secondaryTensor:secondaryTensor name:nil];
-    else
-      return [mpsGraph subtractionWithPrimaryTensor:primaryCastTensor secondaryTensor:secondaryTensor name:nil];
-  };
-  // add alpha's type to the key only if multiply was added to graph
-  binaryOpTensor(self,
-                 other,
-                 alpha,
-                 output,
-                 op_name + "_out_mps:" + (alpha_has_value ? getMPSTypeString(alpha.type()) : ""),
-                 add_sub_lerp_op_block);
 }
 
 } // namespace mps
 
-#define CREATE_MPS_BINARY_COMPARISON_OP_FUNC(func_out, func_stub, other_type)                       \
-  Tensor& func_out(const Tensor& self, const other_type& other, Tensor& output) {                   \
-    mps::binaryOp##other_type(                                                                      \
-        self,                                                                                       \
-        other,                                                                                      \
-        Scalar(1.0),                                                                                \
-        output,                                                                                     \
-        #func_stub,                                                                                 \
-        ^BinaryOpFn(cachedGraph, primaryCastTensor, secondaryCastTensor) {                          \
-          MPSGraph* mpsGraph = cachedGraph->graph();                                                \
-          return [mpsGraph func_stub##                                                              \
-              WithPrimaryTensor:mps::castMPSTensor(mpsGraph, primaryCastTensor, ScalarType::Bool)   \
-                secondaryTensor:mps::castMPSTensor(mpsGraph, secondaryCastTensor, ScalarType::Bool) \
-                           name:nil];                                                               \
-        });                                                                                         \
-    return output;                                                                                  \
+#define CREATE_MPS_BINARY_COMPARISON_OP_FUNC(func_out, func_stub, other_type)                               \
+  Tensor& func_out(const Tensor& self, const other_type& other, Tensor& output) {                           \
+    mps::binaryOp##other_type(                                                                              \
+        self, other, output, #func_stub, ^BinaryOpFn(cachedGraph, primaryCastTensor, secondaryCastTensor) { \
+          MPSGraph* mpsGraph = cachedGraph->graph();                                                        \
+          return [mpsGraph func_stub##                                                                      \
+              WithPrimaryTensor:mps::castMPSTensor(mpsGraph, primaryCastTensor, ScalarType::Bool)           \
+                secondaryTensor:mps::castMPSTensor(mpsGraph, secondaryCastTensor, ScalarType::Bool)         \
+                           name:nil];                                                                       \
+        });                                                                                                 \
+    return output;                                                                                          \
   }
 
-#define CREATE_MPS_STRUCTURED_BINARY_OP_FUNC(func_out, func_stub, other_type)                     \
-  TORCH_IMPL_FUNC(func_out)(const Tensor& self, const other_type& other, const Tensor& output) {  \
-    mps::binaryOp##other_type(self,                                                               \
-                              other,                                                              \
-                              Scalar(1.0),                                                        \
-                              output,                                                             \
-                              #func_stub,                                                         \
-                              ^BinaryOpFn(cachedGraph, primaryCastTensor, secondaryCastTensor) {  \
-                                MPSGraph* mpsGraph = cachedGraph->graph();                        \
-                                return [mpsGraph func_stub##WithPrimaryTensor:primaryCastTensor   \
-                                                              secondaryTensor:secondaryCastTensor \
-                                                                         name:nil];               \
-                              });                                                                 \
+#define CREATE_MPS_STRUCTURED_BINARY_OP_FUNC(func_out, func_stub, other_type)                               \
+  TORCH_IMPL_FUNC(func_out)(const Tensor& self, const other_type& other, const Tensor& output) {            \
+    mps::binaryOp##other_type(                                                                              \
+        self, other, output, #func_stub, ^BinaryOpFn(cachedGraph, primaryCastTensor, secondaryCastTensor) { \
+          MPSGraph* mpsGraph = cachedGraph->graph();                                                        \
+          return [mpsGraph func_stub##WithPrimaryTensor:primaryCastTensor                                   \
+                                        secondaryTensor:secondaryCastTensor                                 \
+                                                   name:nil];                                               \
+        });                                                                                                 \
   }
 
 // output of Boolean Ops will be cast to "MPSDataTypeBool" at the end of binaryOpTensor()
-#define CREATE_MPS_STRUCTURED_BOOLEAN_OP_FUNC(func_out, func_stub, other_type)                    \
-  TORCH_IMPL_FUNC(func_out)(const Tensor& self, const other_type& other, const Tensor& output) {  \
-    mps::binaryOp##other_type(self,                                                               \
-                              other,                                                              \
-                              Scalar(1.0),                                                        \
-                              output,                                                             \
-                              #func_stub,                                                         \
-                              ^BinaryOpFn(cachedGraph, primaryCastTensor, secondaryCastTensor) {  \
-                                MPSGraph* mpsGraph = cachedGraph->graph();                        \
-                                return [mpsGraph func_stub##WithPrimaryTensor:primaryCastTensor   \
-                                                              secondaryTensor:secondaryCastTensor \
-                                                                         name:nil];               \
-                              });                                                                 \
+#define CREATE_MPS_STRUCTURED_BOOLEAN_OP_FUNC(func_out, func_stub, other_type)                              \
+  TORCH_IMPL_FUNC(func_out)(const Tensor& self, const other_type& other, const Tensor& output) {            \
+    mps::binaryOp##other_type(                                                                              \
+        self, other, output, #func_stub, ^BinaryOpFn(cachedGraph, primaryCastTensor, secondaryCastTensor) { \
+          MPSGraph* mpsGraph = cachedGraph->graph();                                                        \
+          return [mpsGraph func_stub##WithPrimaryTensor:primaryCastTensor                                   \
+                                        secondaryTensor:secondaryCastTensor                                 \
+                                                   name:nil];                                               \
+        });                                                                                                 \
   }
 
 // Boolean Binary Ops
@@ -395,20 +335,16 @@ static void add_sub_lerp_template(const Tensor& self,
   if (!mps::supportsComplex() && (c10::isComplexType(self.scalar_type()) || c10::isComplexType(other.scalar_type()))) {
     return mps::complex_mul_out(self, other, output);
   }
-  mps::binaryOpTensor(
-      self, other, Scalar(1.0), output, "mul", ^BinaryOpFn(cachedGraph, primaryCastTensor, secondaryCastTensor) {
-        MPSGraph* mpsGraph = cachedGraph->graph();
-        return [mpsGraph multiplicationWithPrimaryTensor:primaryCastTensor
-                                         secondaryTensor:secondaryCastTensor
-                                                    name:nil];
-      });
+  mps::binaryOpTensor(self, other, output, "mul", ^BinaryOpFn(cachedGraph, primaryCastTensor, secondaryCastTensor) {
+    MPSGraph* mpsGraph = cachedGraph->graph();
+    return [mpsGraph multiplicationWithPrimaryTensor:primaryCastTensor secondaryTensor:secondaryCastTensor name:nil];
+  });
 }
 TORCH_IMPL_FUNC(atan2_out_mps)(const Tensor& self, const Tensor& other, const Tensor& output) {
-  mps::binaryOpTensor(
-      self, other, Scalar(1.0), output, "atan2", ^BinaryOpFn(cachedGraph, primaryCastTensor, secondaryCastTensor) {
-        MPSGraph* mpsGraph = cachedGraph->graph();
-        return [mpsGraph atan2WithPrimaryTensor:primaryCastTensor secondaryTensor:secondaryCastTensor name:nil];
-      });
+  mps::binaryOpTensor(self, other, output, "atan2", ^BinaryOpFn(cachedGraph, primaryCastTensor, secondaryCastTensor) {
+    MPSGraph* mpsGraph = cachedGraph->graph();
+    return [mpsGraph atan2WithPrimaryTensor:primaryCastTensor secondaryTensor:secondaryCastTensor name:nil];
+  });
 }
 
 TORCH_IMPL_FUNC(div_out_mode_mps)
@@ -421,28 +357,10 @@ static void add_sub_lerp_template(const Tensor& self,
 }
 
 TORCH_IMPL_FUNC(add_out_mps)(const Tensor& self, const Tensor& other, const Scalar& alpha, const Tensor& output) {
-  if ((isComplexType(self.scalar_type()) || isComplexType(other.scalar_type())) && !alpha.isComplex() &&
-      !mps::supportsComplex()) {
-    // Complex add with non-complex alpha is just add over views
-    return mps::add_sub_lerp_template(mps::legacy_complex_as_view(self),
-                                      mps::legacy_complex_as_view(other),
-                                      alpha,
-                                      mps::legacy_complex_as_view(output),
-                                      "add");
-  }
   mps::add_sub_lerp_template(self, other, alpha, output, "add");
 }
 
 TORCH_IMPL_FUNC(sub_out_mps)(const Tensor& self, const Tensor& other, const Scalar& alpha, const Tensor& output) {
-  if ((isComplexType(self.scalar_type()) || isComplexType(other.scalar_type())) && !alpha.isComplex() &&
-      !mps::supportsComplex()) {
-    // Complex sub with non-complex alpha is just add over views
-    return mps::add_sub_lerp_template(mps::legacy_complex_as_view(self),
-                                      mps::legacy_complex_as_view(other),
-                                      alpha,
-                                      mps::legacy_complex_as_view(output),
-                                      "sub");
-  }
   mps::add_sub_lerp_template(self, other, alpha, output, "sub");
 }
 
@@ -479,7 +397,7 @@ static void div_floor_kernel_mps(TensorIteratorBase& iter) {
                                                                name:nil];
     return [mpsGraph squareRootWithTensor:sumTensor name:nil];
   };
-  mps::binaryOpTensor(self, other, Scalar(1.0), output, "hypot_out_mps", hypot_op_block);
+  mps::binaryOpTensor(self, other, output, "hypot_out_mps", hypot_op_block);
 }
 
 TORCH_IMPL_FUNC(logaddexp_out_mps)(const Tensor& self, const Tensor& other, const Tensor& output) {
@@ -491,7 +409,7 @@ static void div_floor_kernel_mps(TensorIteratorBase& iter) {
                                        name:nil];
     return [mpsGraph logarithmWithTensor:sumTensor name:nil];
   };
-  mps::binaryOpTensor(self, other, Scalar(1.0), output, "logaddexp_out_mps", logaddexp_op_block);
+  mps::binaryOpTensor(self, other, output, "logaddexp_out_mps", logaddexp_op_block);
 }
 
 TORCH_IMPL_FUNC(logaddexp2_out_mps)(const Tensor& self, const Tensor& other, const Tensor& output) {
@@ -503,7 +421,7 @@ static void div_floor_kernel_mps(TensorIteratorBase& iter) {
                                        name:nil];
     return [mpsGraph logarithmBase2WithTensor:sumTensor name:nil];
   };
-  mps::binaryOpTensor(self, other, Scalar(1.0), output, "logaddexp2_out_mps", logaddexp2_op_block);
+  mps::binaryOpTensor(self, other, output, "logaddexp2_out_mps", logaddexp2_op_block);
 }
 
 TORCH_IMPL_FUNC(xlogy_out_mps)(const Tensor& self, const Tensor& other, const Tensor& output) {
@@ -528,7 +446,7 @@ static void div_floor_kernel_mps(TensorIteratorBase& iter) {
                                                   name:nil];
     return outputTensor;
   };
-  mps::binaryOpTensor(self, other, Scalar(1.0), output, "xlogy_out_mps", xlogy_op_block);
+  mps::binaryOpTensor(self, other, output, "xlogy_out_mps", xlogy_op_block);
 }
 
 TORCH_IMPL_FUNC(lerp_Scalar_mps)(const Tensor& self, const Tensor& end, const Scalar& weight, const Tensor& out) {
