[MPS][BE] Add copysign integral flavors as functor (pytorch#147183)

Pull Request resolved: pytorch#147183
Approved by: https://github.com/dcci
ghstack dependencies: pytorch#147182

Author: malfet

aten/src/ATen/native/mps/kernels/BinaryKernel.metal

@@ -18,9 +18,17 @@ struct fmin_functor {
 
 struct copysign_functor {
   template <typename T>
-  inline T operator()(const T a, const T b) {
+  inline enable_if_t<is_floating_point_v<T>, T> operator()(
+      const T a,
+      const T b) {
     return static_cast<T>(::metal::copysign(a, b));
   }
+  template <typename T>
+  inline enable_if_t<!is_floating_point_v<T>, float> operator()(
+      const T a,
+      const T b) {
+    return ::metal::copysign(static_cast<float>(a), static_cast<float>(b));
+  }
 };
 
 struct zeta_functor {
@@ -111,20 +119,6 @@ kernel void binary_dense(
       device result_of<DTYPE, NAME##_functor> * out_,        \
       uint tid)
 
-template <typename T>
-kernel void copysign_integral(
-    constant void* input_ [[buffer(0)]],
-    constant void* other_ [[buffer(1)]],
-    device void* out_ [[buffer(2)]],
-    constant uint3* offsets [[buffer(3)]],
-    uint tid [[thread_position_in_grid]]) {
-  device float* out = (device float*)((device uint8_t*)out_ + offsets[tid].x);
-  constant T* input = (constant T*)((constant uint8_t*)input_ + offsets[tid].y);
-  constant T* other = (constant T*)((constant uint8_t*)other_ + offsets[tid].z);
-
-  *out = copysign(static_cast<float>(*input), static_cast<float>(*other));
-}
-
 #define REGISTER_BINARY_OP(NAME, DTYPE)                             \
   template [[host_name(#NAME "_" #DTYPE)]] kernel void NAME<DTYPE>( \
       constant void* input_,                                        \
@@ -133,17 +127,14 @@ kernel void copysign_integral(
       constant uint3* offsets,                                      \
       uint tid)
 
-#define REGISTER_COPYSIGN_INTEGRAL_OP(DTYPE)             \
-  template [[host_name("copysign_" #DTYPE)]] kernel void \
-  copysign_integral<DTYPE>(                              \
-      constant void* input_ [[buffer(0)]],               \
-      constant void* other_ [[buffer(1)]],               \
-      device void* out_ [[buffer(2)]],                   \
-      constant uint3* offsets [[buffer(3)]],             \
-      uint tid [[thread_position_in_grid]]);
-
+REGISTER_BINARY_INDEXING_OP(copysign, long);
+REGISTER_BINARY_INDEXING_OP(copysign, int);
 REGISTER_BINARY_INDEXING_OP(copysign, float);
 REGISTER_BINARY_INDEXING_OP(copysign, half);
+REGISTER_BINARY_INDEXING_OP(copysign, short);
+REGISTER_BINARY_INDEXING_OP(copysign, uchar);
+REGISTER_BINARY_INDEXING_OP(copysign, char);
+REGISTER_BINARY_INDEXING_OP(copysign, bool);
 REGISTER_BINARY_INDEXING_OP(fmax, float);
 REGISTER_BINARY_INDEXING_OP(fmax, half);
 REGISTER_BINARY_INDEXING_OP(fmin, float);
@@ -160,12 +151,6 @@ REGISTER_BINARY_INDEXING_OP(fmin, bfloat);
 REGISTER_BINARY_INDEXING_OP(nextafter, bfloat);
 REGISTER_BINARY_INDEXING_OP(zeta, bfloat);
 #endif
-REGISTER_COPYSIGN_INTEGRAL_OP(int);
-REGISTER_COPYSIGN_INTEGRAL_OP(long);
-REGISTER_COPYSIGN_INTEGRAL_OP(short);
-REGISTER_COPYSIGN_INTEGRAL_OP(char);
-REGISTER_COPYSIGN_INTEGRAL_OP(uchar);
-REGISTER_COPYSIGN_INTEGRAL_OP(bool);
 
 // Complex binary functions
 template <typename T>

aten/src/ATen/native/mps/operations/BinaryKernel.mm

@@ -110,7 +110,7 @@ static void fmin_mps_kernel(TensorIteratorBase& iter) {
 }
 
 static void copysign_mps_kernel(TensorIteratorBase& iter) {
-  mps::binary_mps_impl(iter, "copysign", false);
+  mps::binary_mps_impl(iter, "copysign");
 }
 
 static void nextafter_mps_kernel(TensorIteratorBase& iter) {