[jiterator] abs : complex

As per title

Tests exist in `test_unary_ufuncs` and `test_ops`
Pull Request resolved: pytorch#74455
Approved by: https://github.com/anjali411

aten/src/ATen/native/cuda/AbsKernel.cu

@@ -1,6 +1,7 @@
 #define TORCH_ASSERT_NO_OPERATORS
 #include <ATen/native/UnaryOps.h>
 #include <ATen/native/cuda/Loops.cuh>
+#include <ATen/native/cuda/JitLoops.cuh>
 #include <ATen/Dispatch.h>
 #include <ATen/native/DispatchStub.h>
 #include <ATen/native/TensorIterator.h>
@@ -14,12 +15,36 @@ struct AbsFunctor {
   }
 };
 
+const char abs_name[] = "abs_kernel";
 void abs_kernel_cuda(TensorIteratorBase& iter) {
-  AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND3(ScalarType::Half, ScalarType::BFloat16, ScalarType::Bool, iter.dtype(), "abs_cuda", [&]() {
-    gpu_kernel(iter, AbsFunctor<scalar_t>());
-  });
+  auto dtype = iter.dtype();
+  if (at::isComplexType(dtype)) {
+#if AT_USE_JITERATOR()
+    static const auto abs_string = jiterator_stringify(
+        template <typename T> T abs_kernel(T x) { return std::abs(x); });
+    AT_DISPATCH_COMPLEX_TYPES(dtype, "abs_cuda", [&]() {
+      jitted_gpu_kernel<
+          /*name=*/abs_name,
+          /*return_dtype=*/scalar_t,
+          /*common_dtype=*/scalar_t,
+          /*arity=*/1>(iter, abs_string);
+    });
+#else
+    AT_DISPATCH_COMPLEX_TYPES(dtype, "abs_cuda", [&]() {
+      gpu_kernel(iter, AbsFunctor<scalar_t>());
+    });
+#endif
+  } else {
+    AT_DISPATCH_ALL_TYPES_AND3(
+        ScalarType::Half,
+        ScalarType::BFloat16,
+        ScalarType::Bool,
+        iter.dtype(),
+        "abs_cuda",
+        [&]() { gpu_kernel(iter, AbsFunctor<scalar_t>()); });
+  }
 }
 
-REGISTER_DISPATCH(abs_stub, &abs_kernel_cuda);
+  REGISTER_DISPATCH(abs_stub, &abs_kernel_cuda);
 
 }} // namespace at::native