[jiterator] sigmoid_backward: complex (pytorch#74948)

Summary:
Follows: pytorch#74748

cc kshitij12345!

Pull Request resolved: pytorch#74948

Reviewed By: malfet

Differential Revision: D35445949

Pulled By: ngimel

fbshipit-source-id: acd921a210fd15cc78046e37532173162daa38b8
(cherry picked from commit 7883b8d)

Author: khushi-411

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

@@ -8,22 +8,41 @@
 #include <ATen/native/DispatchStub.h>
 #include <ATen/native/TensorIterator.h>
 #include <ATen/native/cuda/Loops.cuh>
+#include <ATen/native/cuda/JitLoops.cuh>
 
 // NOTE: CUDA on Windows requires that the enclosing function
 // of a __device__ lambda not have internal linkage.
 
 namespace at {
 namespace native {
 
+const char sigmoid_backward_name[] = "sigmoid_backward";
 void sigmoid_backward_kernel_cuda(TensorIteratorBase& iter) {
-  if(isComplexType(iter.dtype())) {
-    AT_DISPATCH_COMPLEX_TYPES(iter.dtype(), "sigmoid_backward_cuda", [&]() {
+  auto dtype = iter.dtype();
+  if(isComplexType(dtype)) {
+#if AT_USE_JITERATOR()
+    static const auto sigmoid_backward_string = jiterator_stringify(
+        template <typename T>
+        T sigmoid_backward(T a, T b) {
+          return a * std::conj((T{1.} - b) * b);
+        }
+    ); // sigmoid_backward_string
+    AT_DISPATCH_COMPLEX_TYPES(dtype, "sigmoid_backward_cuda", [&]() {
+        jitted_gpu_kernel<
+          /*name=*/ sigmoid_backward_name,
+          /*return_dtype=*/ scalar_t,
+          /*common_dtype=*/ scalar_t,
+          /*arity=*/ 2>(iter, sigmoid_backward_string);
+    });
+#else
+    AT_DISPATCH_COMPLEX_TYPES(dtype, "sigmoid_backward_cuda", [&]() {
       gpu_kernel(iter, [] GPU_LAMBDA(scalar_t a, scalar_t b) -> scalar_t {
         return a * std::conj((scalar_t{1.} - b) * b);
       });
     });
+#endif
   } else {
-    AT_DISPATCH_FLOATING_TYPES_AND2(at::ScalarType::Half, at::ScalarType::BFloat16, iter.dtype(), "sigmoid_backward_cuda", [&]() {
+    AT_DISPATCH_FLOATING_TYPES_AND2(at::ScalarType::Half, at::ScalarType::BFloat16, dtype, "sigmoid_backward_cuda", [&]() {
       gpu_kernel(iter, []GPU_LAMBDA(scalar_t a, scalar_t b) -> scalar_t {
         return a * (scalar_t(1.) - b) * b;
       });