[MPS] Add support for modified_bessel_k1 to eager and inductor. (pytorch#149687)

Pull Request resolved: pytorch#149687
Approved by: https://github.com/malfet

Author: dcci

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

@@ -8,6 +8,7 @@ DEFINE_UNARY_FLOATING_FUNCTOR(bessel_j1_forward);
 DEFINE_UNARY_FLOATING_FUNCTOR(modified_bessel_i0_forward);
 DEFINE_UNARY_FLOATING_FUNCTOR(modified_bessel_i1_forward);
 DEFINE_UNARY_FLOATING_FUNCTOR(modified_bessel_k0_forward);
+DEFINE_UNARY_FLOATING_FUNCTOR(modified_bessel_k1_forward);
 DEFINE_UNARY_FLOATING_FUNCTOR(i0);
 DEFINE_UNARY_FLOATING_FUNCTOR(i0e);
 DEFINE_UNARY_FLOATING_FUNCTOR(i1);
@@ -53,6 +54,7 @@ struct bessel_y1_forward_functor {
   REGISTER_UNARY_OP(modified_bessel_i0_forward, DTI, DTO); \
   REGISTER_UNARY_OP(modified_bessel_i1_forward, DTI, DTO); \
   REGISTER_UNARY_OP(modified_bessel_k0_forward, DTI, DTO); \
+  REGISTER_UNARY_OP(modified_bessel_k1_forward, DTI, DTO); \
   REGISTER_UNARY_OP(bessel_y0_forward, DTI, DTO);          \
   REGISTER_UNARY_OP(bessel_y1_forward, DTI, DTO);          \
   REGISTER_UNARY_OP(i0, DTI, DTO);                         \

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

@@ -56,6 +56,10 @@ static void modified_bessel_k0_kernel_mps(TensorIteratorBase& iter) {
   lib.exec_unary_kernel(iter, "modified_bessel_k0_forward");
 }
 
+static void modified_bessel_k1_kernel_mps(TensorIteratorBase& iter) {
+  lib.exec_unary_kernel(iter, "modified_bessel_k1_forward");
+}
+
 static void bessel_y0_kernel_mps(TensorIteratorBase& iter) {
   lib.exec_unary_kernel(iter, "bessel_y0_forward");
 }
@@ -73,6 +77,7 @@ static void bessel_y1_kernel_mps(TensorIteratorBase& iter) {
 REGISTER_DISPATCH(special_modified_bessel_i0_stub, &modified_bessel_i0_kernel_mps)
 REGISTER_DISPATCH(special_modified_bessel_i1_stub, &modified_bessel_i1_kernel_mps)
 REGISTER_DISPATCH(special_modified_bessel_k0_stub, &modified_bessel_k0_kernel_mps)
+REGISTER_DISPATCH(special_modified_bessel_k1_stub, &modified_bessel_k1_kernel_mps)
 REGISTER_DISPATCH(special_bessel_y0_stub, &bessel_y0_kernel_mps)
 REGISTER_DISPATCH(special_bessel_y1_stub, &bessel_y1_kernel_mps)
 REGISTER_DISPATCH(special_spherical_bessel_j0_stub, &spherical_bessel_j0_kernel_mps)

aten/src/ATen/native/native_functions.yaml

@@ -15486,7 +15486,7 @@
 
 - func: special_modified_bessel_k1.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!)
   dispatch:
-    CPU, CUDA: special_modified_bessel_k1_out
+    CPU, CUDA, MPS: special_modified_bessel_k1_out
   python_module: special
   structured_inherits: TensorIteratorBase
   structured: True

c10/metal/special_math.h

@@ -1320,5 +1320,73 @@ inline float modified_bessel_k0_forward(T x) {
   return ::metal::exp(-x) * (0.5 * (b - p)) / ::metal::sqrt(x);
 } // modified_bessel_k0_forward(T x)
 
+template <typename T>
+float modified_bessel_k1_forward(T x) {
+  constexpr float A[] = {
+      -7.02386347938628759343e-18,
+      -2.42744985051936593393e-15,
+      -6.66690169419932900609e-13,
+      -1.41148839263352776110e-10,
+      -2.21338763073472585583e-08,
+      -2.43340614156596823496e-06,
+      -1.73028895751305206302e-04,
+      -6.97572385963986435018e-03,
+      -1.22611180822657148235e-01,
+      -3.53155960776544875667e-01,
+      +1.52530022733894777053e+00,
+  };
+
+  constexpr float B[] = {
+      -5.75674448366501715755e-18, +1.79405087314755922667e-17,
+      -5.68946255844285935196e-17, +1.83809354436663880070e-16,
+      -6.05704724837331885336e-16, +2.03870316562433424052e-15,
+      -7.01983709041831346144e-15, +2.47715442448130437068e-14,
+      -8.97670518232499435011e-14, +3.34841966607842919884e-13,
+      -1.28917396095102890680e-12, +5.13963967348173025100e-12,
+      -2.12996783842756842877e-11, +9.21831518760500529508e-11,
+      -4.19035475934189648750e-10, +2.01504975519703286596e-09,
+      -1.03457624656780970260e-08, +5.74108412545004946722e-08,
+      -3.50196060308781257119e-07, +2.40648494783721712015e-06,
+      -1.93619797416608296024e-05, +1.95215518471351631108e-04,
+      -2.85781685962277938680e-03, +1.03923736576817238437e-01,
+      +2.72062619048444266945e+00,
+  };
+
+  if (x == 0.0) {
+    return INFINITY;
+  }
+
+  if (x < 0.0) {
+    return NAN;
+  }
+
+  float p;
+  float q = 0.0;
+
+  if (x <= 2.0) {
+    float a = A[0];
+
+    for (uint8_t index = 1; index < 11; index++) {
+      p = q;
+      q = a;
+      a = (x * x - T(2.0)) * q - p + A[index];
+    }
+
+    return ::metal::precise::log(T(0.5) * x) * modified_bessel_i1_forward(x) +
+        0.5 * (a - p) / x;
+  }
+
+  float b = B[0];
+
+  for (uint8_t index = 1; index < 25; index++) {
+    p = q;
+    q = b;
+    b = (8.0 / x - 2.0) * q - p + B[index];
+  }
+
+  return ::metal::precise::exp(-x) * (0.5 * (b - p)) /
+      ::metal::precise::sqrt(x);
+}
+
 } // namespace metal
 } // namespace c10

test/inductor/test_mps_basic.py

@@ -98,6 +98,7 @@ def test_cast(self, dtype):
         "modified_bessel_i0",
         "modified_bessel_i1",
         "modified_bessel_k0",
+        "modified_bessel_k1",
         "entr",
     ]
 

test/test_mps.py

@@ -663,7 +663,6 @@ def mps_ops_modifier(ops):
         'special.hermite_polynomial_he': None,
         'special.laguerre_polynomial_l': None,
         'special.log_ndtr': None,
-        'special.modified_bessel_k1': None,
         'special.ndtri': None,
         'special.scaled_modified_bessel_k0': None,
         'special.scaled_modified_bessel_k1': None,

torch/_inductor/codegen/mps.py

@@ -414,6 +414,10 @@ def modified_bessel_i1(x: CSEVariable) -> str:
     def modified_bessel_k0(x: CSEVariable) -> str:
         return f"c10::metal::modified_bessel_k0_forward({x})"
 
+    @staticmethod
+    def modified_bessel_k1(x: CSEVariable) -> str:
+        return f"c10::metal::modified_bessel_k1_forward({x})"
+
 
 MetalOverrides._initialize_pointwise_overrides("mps")