1 #define TORCH_ASSERT_NO_OPERATORS
2 #include <ATen/AccumulateType.h>
3 #include <ATen/Dispatch.h>
4 #include <ATen/native/BinaryOps.h>
5 #include <ATen/native/DispatchStub.h>
6 #include <ATen/native/TensorIterator.h>
7 #include <ATen/native/cuda/BinaryInternal.h>
8 #include <c10/cuda/CUDAGuard.h>
9 #include <c10/cuda/CUDAMathCompat.h>
10 #include <c10/util/TypeSafeSignMath.h>
11 #include <ATen/native/cuda/JitLoops.cuh>
12 #include <ATen/native/cuda/Loops.cuh>
13
14 #include <type_traits>
15
16 // NOTE: CUDA on Windows requires that the enclosing function
17 // of a __device__ lambda not have internal linkage.
18
19 namespace at::native {
20
21 CONSTEXPR_EXCEPT_WIN_CUDA char mul_name[] = "mul_kernel";
mul_kernel_cuda(TensorIteratorBase & iter)22 void mul_kernel_cuda(TensorIteratorBase& iter) {
23 auto common_dtype = iter.common_dtype();
24 if (common_dtype == kComplexHalf) {
25 using scalar_t = c10::complex<at::Half>;
26 #if AT_USE_JITERATOR()
27 static const auto mul_string = jiterator_stringify(
28 template <typename T> T mul_kernel(T a, T b) { return a * b; });
29 opmath_jitted_gpu_kernel_with_scalars<mul_name, scalar_t, scalar_t>(
30 iter, mul_string);
31 #else
32 using opmath_t = at::opmath_type<scalar_t>;
33 opmath_symmetric_gpu_kernel_with_scalars<scalar_t>(
34 iter, binary_internal::MulFunctor<opmath_t>());
35 #endif
36 } else {
37 AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND3(
38 kHalf, kBFloat16, kBool, iter.common_dtype(), "mul_cuda", [&]() {
39 using opmath_t = at::opmath_type<scalar_t>;
40 opmath_symmetric_gpu_kernel_with_scalars<scalar_t>(
41 iter, binary_internal::MulFunctor<opmath_t>());
42 });
43 }
44 }
45
46 REGISTER_DISPATCH(mul_stub, &mul_kernel_cuda);
47
48 } // namespace at::native
49