Add aten::special_airy_ai.out (#1039)

- `special_airy_ai.out`
- `special_airy_ai`

---------

Co-authored-by: Yutao Xu <[email protected]>

Author: min-jean-cho

src/ATen/native/xpu/AiryAi.cpp

@@ -0,0 +1,11 @@
+#include <ATen/native/UnaryOps.h>
+#include <ATen/native/DispatchStub.h>
+#include <ATen/native/TensorIterator.h>
+#include <ATen/native/xpu/sycl/AiryAiKernel.h>
+
+namespace at {
+namespace native {
+REGISTER_XPU_DISPATCH(special_airy_ai_stub, &xpu::airy_ai_kernel);
+
+} // namespace native
+} // namespace at

src/ATen/native/xpu/XPUFallback.template

@@ -200,7 +200,6 @@ TORCH_LIBRARY_IMPL(aten, XPU, m) {
     "segment_reduce",
     "_segment_reduce_backward",
     "sinc.out",
-    "special_airy_ai.out",
     "_thnn_fused_gru_cell",
     "_to_sparse",
     "_to_sparse_csr",

src/ATen/native/xpu/sycl/AiryAiKernel.cpp

@@ -0,0 +1,22 @@
+#include <ATen/ATen.h>
+#include <ATen/native/TensorIterator.h>
+#include <ATen/native/xpu/sycl/Loops.h>
+#include <ATen/native/xpu/sycl/MathExtensions.h>
+
+#include <ATen/native/xpu/sycl/AiryAiKernel.h>
+
+namespace at::native::xpu {
+template <typename scalar_t>
+struct AiryAiFunctor {
+  scalar_t operator()(scalar_t a) const {
+    return airy_ai_forward(a);
+  }
+};
+
+void airy_ai_kernel(TensorIteratorBase& iter) {
+  AT_DISPATCH_FLOATING_TYPES(iter.common_dtype(), "airy_ai_xpu", [&]() {
+    gpu_kernel(iter, AiryAiFunctor<scalar_t>());
+  });
+}
+
+} // namespace at::native::xpu

src/ATen/native/xpu/sycl/AiryAiKernel.h

@@ -0,0 +1,9 @@
+#pragma once
+
+#include <ATen/native/TensorIterator.h>
+
+namespace at::native::xpu {
+
+TORCH_XPU_API void airy_ai_kernel(TensorIteratorBase& iter);
+
+} // namespace at::native::xpu

src/ATen/native/xpu/sycl/MathExtensions.h

@@ -1318,4 +1318,187 @@ static inline C10_HOST_DEVICE scalar_t bessel_y1_forward(scalar_t x) {
       scalar_t(0.797884560802865355879892119868763737) / std::sqrt(x);
 } // bessel_y1_forward(scalar_t x)
 
+template <typename T>
+static inline C10_HOST_DEVICE T airy_ai_forward(T x) {
+  static const T AN[] = {
+      +3.46538101525629032477e-01f,
+      +1.20075952739645805542e+01f,
+      +7.62796053615234516538e+01f,
+      +1.68089224934630576269e+02f,
+      +1.59756391350164413639e+02f,
+      +7.05360906840444183113e+01f,
+      +1.40264691163389668864e+01f,
+      +9.99999999999999995305e-01f,
+  };
+
+  static const T AD[] = {
+      +5.67594532638770212846e-01f,
+      +1.47562562584847203173e+01f,
+      +8.45138970141474626562e+01f,
+      +1.77318088145400459522e+02f,
+      +1.64234692871529701831e+02f,
+      +7.14778400825575695274e+01f,
+      +1.40959135607834029598e+01f,
+      +1.00000000000000000470e+00f,
+  };
+
+  static const T AFN[] = {
+      -1.31696323418331795333e-01f,
+      -6.26456544431912369773e-01f,
+      -6.93158036036933542233e-01f,
+      -2.79779981545119124951e-01f,
+      -4.91900132609500318020e-02f,
+      -4.06265923594885404393e-03f,
+      -1.59276496239262096340e-04f,
+      -2.77649108155232920844e-06f,
+      -1.67787698489114633780e-08f,
+  };
+
+  static const T AFD[] = {
+      +1.33560420706553243746e+01f,
+      +3.26825032795224613948e+01f,
+      +2.67367040941499554804e+01f,
+      +9.18707402907259625840e+00f,
+      +1.47529146771666414581e+00f,
+      +1.15687173795188044134e-01f,
+      +4.40291641615211203805e-03f,
+      +7.54720348287414296618e-05f,
+      +4.51850092970580378464e-07f,
+  };
+
+  static const T AGN[] = {
+      +1.97339932091685679179e-02f,
+      +3.91103029615688277255e-01f,
+      +1.06579897599595591108e+00f,
+      +9.39169229816650230044e-01f,
+      +3.51465656105547619242e-01f,
+      +6.33888919628925490927e-02f,
+      +5.85804113048388458567e-03f,
+      +2.82851600836737019778e-04f,
+      +6.98793669997260967291e-06f,
+      +8.11789239554389293311e-08f,
+      +3.41551784765923618484e-10f,
+  };
+
+  static const T AGD[] = {
+      +9.30892908077441974853e+00f,
+      +1.98352928718312140417e+01f,
+      +1.55646628932864612953e+01f,
+      +5.47686069422975497931e+00f,
+      +9.54293611618961883998e-01f,
+      +8.64580826352392193095e-02f,
+      +4.12656523824222607191e-03f,
+      +1.01259085116509135510e-04f,
+      +1.17166733214413521882e-06f,
+      +4.91834570062930015649e-09f,
+  };
+
+  int domain_flag = 0;
+
+  T ai;
+
+  if (std::isinf(x)) {
+    return std::numeric_limits<T>::quiet_NaN();
+  }
+
+  if (x > T(103.892f)) {
+    return T(0.0f);
+  }
+
+  T f;
+  T g;
+  T k;
+
+  if (x < T(-2.09f)) {
+    T z = T(1.0f) / (T(-2.0f) * x * std::sqrt(-x) / T(3.0f));
+
+    T afn = 0.0f;
+
+    for (uint8_t index = 0; index <= 8; index++) {
+      afn = afn * (z * z) + AFN[index];
+    }
+
+    T afd = 0.0f;
+
+    for (uint8_t index = 0; index <= 8; index++) {
+      afd = afd * (z * z) + AFD[index];
+    }
+
+    T agn = 0.0f;
+
+    for (uint8_t index = 0; index <= 10 + 0; index++) {
+      agn = agn * (z * z) + AGN[index];
+    }
+
+    T agd = 0.0f;
+
+    for (uint8_t index = 0; index <= 10 - 1; index++) {
+      agd = agd * (z * z) + AGD[index];
+    }
+
+    T t = T(-2.0f) * x * std::sqrt(-x) / T(3.0f) +
+        T(0.25f) * T(3.14159265358979323846f);
+
+    return T(5.64189583547756286948e-01f) / std::sqrt(std::sqrt(-x)) *
+        (std::sin(t) * (T(1.0f) + z * z * afn / afd) -
+         std::cos(t) * (z * agn / agd));
+  }
+
+  if (x >= T(2.09f)) {
+    domain_flag = 5;
+
+    T zeta = T(2.0f) * x * std::sqrt(x) / T(3.0f);
+
+    T an = 0.0f;
+
+    for (uint8_t index = 0; index <= 7; index++) {
+      an = an * (T(1.0f) / zeta) + AN[index];
+    }
+
+    T ad = 0.0f;
+
+    for (uint8_t index = 0; index <= 7; index++) {
+      ad = ad * (T(1.0f) / zeta) + AD[index];
+    }
+
+    ai = T(5.64189583547756286948e-01f) * (an / ad) /
+        (T(2.0f) * std::sqrt(std::sqrt(x)) * std::exp(zeta));
+
+    if (x > T(8.3203353f)) {
+      return ai;
+    }
+  }
+
+  f = 1.0f;
+  g = x;
+  k = 1.0f;
+
+  T m = 1.0f;
+  T n = x;
+  T t = 1.0f;
+  T z = x * x * x;
+
+  while (t > std::numeric_limits<T>::epsilon()) {
+    m *= z;
+    k += T(1.0f);
+    m /= k;
+    n *= z;
+    k += T(1.0f);
+    n /= k;
+    m /= k;
+    f += m;
+    k += T(1.0f);
+    n /= k;
+    g += n;
+
+    t = std::abs(m / f);
+  }
+
+  if ((domain_flag & 1) == 0) {
+    return T(0.355028053887817239260f) * f - T(0.258819403792806798405f) * g;
+  }
+
+  return ai;
+} // T airy_ai(T x)
+
 } // namespace at::native::xpu

yaml/native/native_functions.yaml

@@ -6730,6 +6730,21 @@
     XPU: angle_out
   tags: pointwise
 
+- func: special_airy_ai(Tensor x) -> Tensor
+  python_module: special
+  structured_delegate: special_airy_ai.out
+  variants: function
+  tags: pointwise
+
+- func: special_airy_ai.out(Tensor x, *, Tensor(a!) out) -> Tensor(a!)
+  dispatch:
+    XPU: special_airy_ai_out
+  python_module: special
+  structured_inherits: TensorIteratorBase
+  structured: True
+  variants: function
+  tags: pointwise
+
 - func: special_bessel_j0(Tensor self) -> Tensor
   python_module: special
   structured_delegate: special_bessel_j0.out