Revert "Remove conj kernels for real dtypes (pytorch#80374)"

This reverts commit ad44079.

Reverted pytorch#80374 on behalf of https://github.com/atalman due to Breaks internal build UnaryOpsKernel.cpp:208:5: error: unused type alias 'scalar_t' [-Werror,-Wunused-local-typedef]

Author: pytorchmergebot

aten/src/ATen/native/cpu/CopyKernel.cpp

@@ -13,6 +13,9 @@ namespace native {
 inline namespace CPU_CAPABILITY {
 void neg_kernel(TensorIteratorBase &iter);
 void conj_kernel(TensorIteratorBase &iter);
+} // namespace CPU_CAPABILITY
+
+namespace {
 
 void float_bfloat16_copy_kernel(TensorIteratorBase &iter, bool requires_neg) {
   auto strides_out = iter.strides(0);
@@ -269,7 +272,7 @@ void copy_kernel(TensorIterator& iter, bool /*non_blocking*/) {
   }
 }
 
-} // namespace CPU_CAPABILITY
+} // anonymous namespace
 
 REGISTER_DISPATCH(copy_stub, &copy_kernel);
 

aten/src/ATen/native/cpu/CopyKernel.h

@@ -13,7 +13,6 @@
 #include <ATen/cpu/vec/vec.h>
 #include <ATen/cpu/vml.h>
 #include <ATen/native/TensorIterator.h>
-#include <ATen/native/cpu/CopyKernel.h>
 #include <ATen/native/cpu/Loops.h>
 #include <ATen/native/cpu/zmath.h>
 #include <ATen/OpMathType.h>
@@ -204,18 +203,13 @@ static void angle_kernel(TensorIteratorBase& iter) {
 
 // NB: Ignores the negative bit on tensors
 void conj_kernel(TensorIteratorBase& iter) {
-  AT_DISPATCH_SWITCH(iter.common_dtype(), "conj_cpu",
-    AT_DISPATCH_CASE_ALL_TYPES_AND3(kBool, kBFloat16, kHalf, [&] {
-      // conj is a no-op for non-complex types
-      direct_copy_kernel(iter);
-    })
-    AT_DISPATCH_CASE_COMPLEX_TYPES_AND(kComplexHalf, [&] {
-      cpu_kernel_vec(
-          iter,
-          [=](scalar_t a) -> scalar_t { return conj_impl(a); },
-          [=](Vectorized<scalar_t> a) { return a.conj(); });
-    })
-  );
+  AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND4(
+      kBool, kBFloat16, kHalf, kComplexHalf, iter.common_dtype(), "conj_cpu", [&]() {
+        cpu_kernel_vec(
+            iter,
+            [=](scalar_t a) -> scalar_t { return conj_impl(a); },
+            [=](Vectorized<scalar_t> a) { return a.conj(); });
+      });
 }
 
 static void bitwise_not_kernel(TensorIteratorBase& iter) {

aten/src/ATen/native/cpu/UnaryOpsKernel.cpp

@@ -23,6 +23,7 @@ namespace native {
 void neg_kernel_cuda(TensorIteratorBase &iter);
 void conj_kernel_cuda(TensorIteratorBase &iter);
 
+namespace {
 void direct_copy_kernel_cuda(TensorIteratorBase &iter) {
   ScalarType dtype = iter.dtype(0);
   if (isQIntType(dtype)) {
@@ -42,6 +43,7 @@ void neg_conj_kernel_cuda(TensorIteratorBase &iter) {
     gpu_kernel(iter, [] GPU_LAMBDA(scalar_t x) { return -std::conj(x); });
   });
 }
+}  // namespace (anonymous)
 
 using namespace at::cuda;
 

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

@@ -1,7 +1,6 @@
 #define TORCH_ASSERT_NO_OPERATORS
 #include <limits>
 #include <ATen/native/UnaryOps.h>
-#include <ATen/native/cuda/Copy.h>
 #include <ATen/native/cuda/Loops.cuh>
 #include <ATen/native/cuda/JitLoops.cuh>
 #include <ATen/Dispatch.h>
@@ -59,10 +58,22 @@ void angle_kernel_cuda(TensorIteratorBase& iter) {
   }
 }
 
+// We manually overload conj because std::conj does not work types other than c10::complex.
+template<typename scalar_t>
+__host__ __device__ static inline scalar_t conj_wrapper(scalar_t v) {
+  return v;
+}
+
+template<typename T>
+__host__ __device__ static inline c10::complex<T> conj_wrapper(c10::complex<T> v) {
+  return std::conj(v);
+}
+
 // NB: Ignores the negative bit on tensors
 const char conj_name[] = "conj_kernel";
 void conj_kernel_cuda(TensorIteratorBase& iter) {
-  auto conj_chalf = [&] {
+  auto common_dtype = iter.common_dtype();
+  if (common_dtype == kComplexHalf) {
     using scalar_t = c10::complex<at::Half>;
     #if AT_USE_JITERATOR()
       static const auto conj_string = jiterator_stringify(
@@ -74,23 +85,17 @@ void conj_kernel_cuda(TensorIteratorBase& iter) {
       jitted_gpu_kernel<conj_name, scalar_t, scalar_t, 1>(iter, conj_string);
     #else
       gpu_kernel(iter, [] GPU_LAMBDA(scalar_t a) -> scalar_t {
-          return std::conj(a);
+          return conj_wrapper(a);
       });
     #endif
-  };
-
-  AT_DISPATCH_SWITCH(iter.common_dtype(), "conj_cuda",
-    AT_DISPATCH_CASE_ALL_TYPES_AND3(kBool, kBFloat16, kHalf, [&] {
-      // Conj is a no-op for non-complex types
-      direct_copy_kernel_cuda(iter);
-    })
-    AT_DISPATCH_CASE_COMPLEX_TYPES([&] {
-      gpu_kernel(iter, [] GPU_LAMBDA(scalar_t a) -> scalar_t {
-        return std::conj(a);
-      });
-    })
-    AT_DISPATCH_CASE(kComplexHalf, conj_chalf)
-  );
+  } else {
+    AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND3(
+      kBool, kBFloat16, kHalf, iter.common_dtype(), "conj_cuda", [&]() {
+        gpu_kernel(iter, [] GPU_LAMBDA(scalar_t a) -> scalar_t {
+          return conj_wrapper(a);
+        });
+    });
+  }
 }
 
 REGISTER_DISPATCH(angle_stub, &angle_kernel_cuda);