pytorch · khushi-411 · Mar 29, 2022 · Mar 29, 2022 · Mar 30, 2022 · Mar 30, 2022
diff --git a/aten/src/ATen/cuda/llvm_complex.cpp b/aten/src/ATen/cuda/llvm_complex.cpp
@@ -12,6 +12,7 @@
 
 #include <string>
 #include <ATen/cuda/llvm_jit_strings.h>
+#include <limits>
 
 
 namespace at {
@@ -773,6 +774,37 @@ log2(const complex<_Tp>& __x)
     return log(__x) / log(_Tp(2));
 }
 
+// reciprocal
+
+template<class _Tp>
+inline
+complex<_Tp>
+reciprocal(const complex<_Tp>& __x)
+{
+    // Handle extreme cases for numpy compatibility
+    auto both_inf = [](_Tp real, _Tp imag) {
+        return isinf(real) && isinf(imag);
+    };
+
+    auto either_inf = [](_Tp real, _Tp imag) {
+        return isinf(real) || isinf(imag);
+    };
+
+    auto either_nan = [](_Tp real, _Tp imag) {
+        return isnan(real) || isnan(imag);
+    };
+
+    if (either_nan(__x.real(), __x.imag()) || both_inf(__x.real(), __x.imag())) {
+        // If either is Nan or both are infinite, return {nan, nan}
+        return {std::numeric_limits<_Tp>::quiet_NaN(), std::numeric_limits<_Tp>::quiet_NaN()};
+    } else if (either_inf(__x.real(), __x.imag())) {
+        // If either is Inf, return {0, 0}
+        return {0, 0};
+    }
+    const complex<_Tp> one = complex<_Tp>(1.0, 0);
+    return one/__x;
+}
+
 // sqrt
 
 template<class _Tp>

diff --git a/aten/src/ATen/native/cuda/UnaryFractionKernels.cu b/aten/src/ATen/native/cuda/UnaryFractionKernels.cu
@@ -2,6 +2,7 @@
 #include <limits>
 #include <ATen/native/UnaryOps.h>
 #include <ATen/native/cuda/Loops.cuh>
+#include <ATen/native/cuda/JitLoops.cuh>
 #include <ATen/AccumulateType.h>
 #include <ATen/Dispatch.h>
 #include <ATen/native/DispatchStub.h>
@@ -66,12 +67,12 @@ void floor_kernel_cuda(TensorIteratorBase& iter) {
 }
 
 template <typename scalar_t>
-__host__ __device__ static inline scalar_t reciprocal_wrapper(scalar_t a) {
+C10_HOST_DEVICE static inline scalar_t reciprocal_wrapper(scalar_t a) {
   return static_cast<scalar_t>(1)/a;
 }
 
 template<typename T>
-__host__ __device__ static inline c10::complex<T> reciprocal_wrapper(c10::complex<T> v) {
+C10_HOST_DEVICE static inline c10::complex<T> reciprocal_wrapper(c10::complex<T> v) {
   // Handle extreme cases for numpy compatibility
   auto both_inf = [](T real, T imag) {
     return (::isinf(real) && ::isinf(imag));
@@ -96,15 +97,62 @@ __host__ __device__ static inline c10::complex<T> reciprocal_wrapper(c10::comple
   return one/v;
 }
 
+const char reciprocal_name[] = "reciprocal_kernel";
 void reciprocal_kernel_cuda(TensorIteratorBase& iter) {
-  AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES_AND2(
+  auto dtype = iter.common_dtype();
+  if (at::isComplexType(dtype)) {
+#if AT_USE_JITERATOR()
+  static const auto reciprocal_string = jiterator_stringify(
+      template <typename T>
+      T reciprocal_kernel(T v) {
+        // Handle extreme cases for numpy compatibility
+        auto both_inf = [](T real, T imag) {
+          return std::isinf(real) && std::isinf(imag);
+        };
+
+        auto either_inf = [](T real, T imag) {
+          return std::isinf(real) || std::isinf(imag);
+        };
+
+        auto either_nan = [](T real, T imag) {
+          return std::isnan(real) || std::isnan(imag);
+        };
+
+        if (either_nan(v.real(), v.imag()) || both_inf(v.real(), v.imag())) {
+          // If either is Nan or both are infinite, return {nan, nan}
+          return {std::numeric_limits<T>::quiet_NaN(), std::numeric_limits<T>::quiet_NaN()};
+        } else if (either_inf(v.real(), v.imag())) {
+          // If either is Inf, return {0, 0}
+          return {0, 0};
+        }
+        const c10::complex<T> one = c10::complex<T>(1.0, 0);
+        return one/v;
+      }
+  ); // reciprocal_string
+  AT_DISPATCH_COMPLEX_TYPES(dtype, "reciprocal_cuda", [&]() {
+      jitted_gpu_kernel<
+        /*name=*/ reciprocal_name,
+        /*return_dtype=*/ scalar_t,
+        /*common_dtype=*/ scalar_t,
+        /*arity=*/ 1>(iter, reciprocal_string);
+  });
+#else
+  AT_DISPATCH_COMPLEX_TYPES(dtype, "reciprocal_cuda", [&]() {
+    gpu_kernel(iter, []GPU_LAMBDA(scalar_t a) -> scalar_t {
+      return reciprocal_wrapper(a);
+    });
+  });
+#endif
+  } else {
+    AT_DISPATCH_FLOATING_TYPES_AND2(
       ScalarType::Half, ScalarType::BFloat16,
-      iter.common_dtype(), "reciprocal_cuda",
+      dtype, "reciprocal_cuda",
       [&]() {
         gpu_kernel(iter, []GPU_LAMBDA(scalar_t a) -> scalar_t {
           return reciprocal_wrapper(a);
         });
       });
+  }
 }
 
 // We manually overload nearbyint because std::nearbyint does not work with std::complex types and ROCm.