[MPS] Fix abs for complex types (#125662)

malfet · pytorchmergebot · commit 0fd1fc17c3a5 · 2024-05-07T22:15:20.000Z
By calling `realPartOfTensor:` if input type is complex on Sonoma and fall back to `at::view_as_real` trick on Ventura. Split `unary_op` template into `unary_op` and `unary_op_noresize`, which skips resize and empty checks Marked `abs`, `isclose` and `nn.functional.softsign` OpInfo tests as supported by complex types Fixes #125135 Pull Request resolved: #125662 Approved by: https://github.com/kulinseth
diff --git a/aten/src/ATen/native/mps/MPSGraphSonomaOps.h b/aten/src/ATen/native/mps/MPSGraphSonomaOps.h
@@ -25,6 +25,10 @@ typedef NS_ENUM(NSUInteger, MPSGraphFFTScalingMode)
 -(MPSGraphTensor * _Nonnull) conjugateWithTensor:(MPSGraphTensor * _Nonnull) tensor
                                             name:(NSString * _Nullable) name;
 
+-(MPSGraphTensor * _Nonnull) realPartOfTensor:(MPSGraphTensor * _Nonnull) tensor
+                                         name:(NSString * _Nullable) name;
+
+
 -(MPSGraphTensor * _Nonnull) fastFourierTransformWithTensor:(MPSGraphTensor * _Nonnull) tensor
                                                        axes:(NSArray<NSNumber *> * _Nonnull) axes
                                                  descriptor:(MPSGraphFFTDescriptor * _Nonnull) descriptor
diff --git a/aten/src/ATen/native/mps/operations/UnaryOps.mm b/aten/src/ATen/native/mps/operations/UnaryOps.mm
@@ -75,23 +75,10 @@ static bool is_empty_tensor(const Tensor& self) {
   return self.numel() == 0;
 }
 
-static void unary_op(const Tensor& self,
-                     const Tensor& output_,
-                     std::string op_name,
-                     UnaryOpBlock unaryBlock,
-                     is_noop_p is_noop = is_empty_tensor) {
+static void unary_op_noresize(const Tensor& self, const Tensor& output_, std::string op_name, UnaryOpBlock unaryBlock) {
   TORCH_CHECK(!(!is_macos_13_or_newer() && self.scalar_type() == ScalarType::Byte),
               "MPS support unary op with uint8 natively starting from macOS 13.0");
 
-  if (!output_.is_same_size(self)) {
-    output_.resize_(self.sizes());
-  }
-
-  if (is_noop(self)) {
-    output_.copy_(self);
-    return;
-  }
-
   auto output = output_;
   bool needsCopyToOutput = false;
   if (output.storage_offset() || !output.is_contiguous()) {
@@ -139,6 +126,23 @@ static void unary_op(const Tensor& self,
   }
 }
 
+static void unary_op(const Tensor& self,
+                     const Tensor& output_,
+                     std::string op_name,
+                     UnaryOpBlock unaryBlock,
+                     is_noop_p is_noop = is_empty_tensor) {
+  if (!output_.is_same_size(self)) {
+    output_.resize_(self.sizes());
+  }
+
+  if (is_noop(self)) {
+    output_.copy_(self);
+    return;
+  }
+
+  unary_op_noresize(self, output_, op_name, unaryBlock);
+}
+
 MPSGraphTensor* trunc_tensor(MPSGraph* mpsGraph, MPSGraphTensor* inputTensor) {
   // Rounding is a no-op for integral types, and also a reasonable workaround
   // For MPSGraph bug on Apple Silicon, that throws `Function floorOp_i64 was not found in the library`
@@ -168,6 +172,12 @@ static void unary_op(const Tensor& self,
   return [mpsGraph logarithmWithTensor:addedTensor name:nil];
 }
 
+static MPSGraphTensor* lengthOfComplexAsReal(MPSGraph* mpsGraph, MPSGraphTensor* inputTensor) {
+  auto squares = [mpsGraph squareWithTensor:inputTensor name:nil];
+  auto sumSquares = [mpsGraph reductionSumWithTensor:squares axis:-1 name:nil];
+  return [mpsGraph squareRootWithTensor:sumSquares name:nil];
+}
+
 } // namespace mps
 
 TORCH_IMPL_FUNC(trunc_out_mps)(const Tensor& self, const Tensor& output) {
@@ -226,14 +236,6 @@ static void unary_op(const Tensor& self,
     });                                                                                                          \
   }
 
-#define CREATE_MPS_UNARY_TORCH_IMPL_FUNC(func_out, func_stub)                                                    \
-  Tensor& func_out(const Tensor& self, Tensor& output) {                                                         \
-    mps::unary_op(self, output, #func_out, ^MPSGraphTensor*(MPSGraph * mpsGraph, MPSGraphTensor * inputTensor) { \
-      return [mpsGraph func_stub##WithTensor:inputTensor name:nil];                                              \
-    });                                                                                                          \
-    return output;                                                                                               \
-  }
-
 CREATE_MPS_STRUCTURED_UNARY_TORCH_IMPL_FUNC(exp_out_mps, exponent)
 CREATE_MPS_STRUCTURED_UNARY_TORCH_IMPL_FUNC(exp2_out_mps, exponentBase2)
 CREATE_MPS_STRUCTURED_UNARY_TORCH_IMPL_FUNC(reciprocal_out_mps, reciprocal)
@@ -257,7 +259,35 @@ static void unary_op(const Tensor& self,
 CREATE_MPS_STRUCTURED_UNARY_TORCH_IMPL_FUNC(acosh_out_mps, acosh)
 CREATE_MPS_STRUCTURED_UNARY_TORCH_IMPL_FUNC(atanh_out_mps, atanh)
 
-CREATE_MPS_UNARY_TORCH_IMPL_FUNC(abs_out_mps, absolute)
+Tensor& abs_out_mps(const Tensor& self, Tensor& output) {
+  using namespace mps;
+
+  if (!output.is_same_size(self)) {
+    output.resize_(self.sizes());
+  }
+
+  if (self.numel() == 0) {
+    return output;
+  }
+
+  if (supportsComplex() || !self.is_complex()) {
+    unary_op_noresize(self, output, "abs_out_mps", ^MPSGraphTensor*(MPSGraph* mpsGraph, MPSGraphTensor* inputTensor) {
+      auto rc = [mpsGraph absoluteWithTensor:inputTensor name:nil];
+      if (self.is_complex()) {
+        rc = [mpsGraph realPartOfTensor:rc name:nil];
+      }
+      return rc;
+    });
+  } else {
+    Tensor realInput = at::view_as_real(self);
+    unary_op_noresize(
+        realInput, output, "abs_out_mps", ^MPSGraphTensor*(MPSGraph* mpsGraph, MPSGraphTensor* inputTensor) {
+          auto rc = lengthOfComplexAsReal(mpsGraph, inputTensor);
+          return [mpsGraph reshapeTensor:rc withShape:getMPSShape(output) name:nil];
+        });
+  }
+  return output;
+}
 
 Tensor& logical_not_out_mps(const Tensor& self, Tensor& output) {
   auto bool_self = self.to(ScalarType::Bool);
@@ -484,9 +514,7 @@ static void cumulative_op_impl(const Tensor& self,
   Tensor realOutput = at::view_as_real(output);
 
   auto complex_sgn_op = [&](MPSGraph* mpsGraph, MPSGraphTensor* inputTensor) -> MPSGraphTensor* {
-    MPSGraphTensor* squares = [mpsGraph squareWithTensor:inputTensor name:nil];
-    MPSGraphTensor* sumSquares = [mpsGraph reductionSumWithTensor:squares axis:-1 name:nil];
-    MPSGraphTensor* norm = [mpsGraph squareRootWithTensor:sumSquares name:nil];
+    MPSGraphTensor* norm = mps::lengthOfComplexAsReal(mpsGraph, inputTensor);
     MPSGraphTensor* zero = [mpsGraph constantWithScalar:0.0 dataType:norm.dataType];
     MPSGraphTensor* isZero = [mpsGraph equalWithPrimaryTensor:norm secondaryTensor:zero name:nil];
     MPSGraphTensor* sgnTensor = [mpsGraph divisionWithPrimaryTensor:inputTensor secondaryTensor:norm name:nil];
diff --git a/test/test_mps.py b/test/test_mps.py
@@ -227,6 +227,7 @@ def mps_ops_modifier(ops):
         '__radd__',
         '__rmul__',
         '__getitem__',
+        'abs',
         'add',
         'argwhere',
         'atleast_1d',
@@ -286,8 +287,8 @@ def mps_ops_modifier(ops):
         'narrow_copy',
         'nn.functional.conv1d',
         'nn.functional.conv_transpose1d',
-        'nn.functional.padcircular',
         'nn.functional.feature_alpha_dropoutwithout_train',
+        'nn.functional.padcircular',
         'nn.functional.unfold',
         'nonzero',
         'ones',
@@ -383,6 +384,7 @@ def mps_ops_modifier(ops):
         'half',
         'hstack',
         'int',
+        'isclose',
         'isnan',
         'ldexp',
         'log10',
@@ -403,12 +405,13 @@ def mps_ops_modifier(ops):
         'mean',
         'ne',
         'neg',
-        'nn.functional.rms_norm',
         'nn.functional.padconstant',
         'nn.functional.padreflect',
         'nn.functional.padreplicate',
         'nn.functional.pixel_shuffle',
         'nn.functional.pixel_unshuffle',
+        'nn.functional.rms_norm',
+        'nn.functional.softsign',
         'nn.functional.tanhshrink',
         'prod',
         'reciprocal',
