use opmath_t in jiterator

khushi-411 · khushi-411 · commit bf94d1016fb3 · 2022-04-20T10:50:53.000+05:30
diff --git a/aten/src/ATen/cuda/llvm_complex.cpp b/aten/src/ATen/cuda/llvm_complex.cpp
@@ -874,7 +874,7 @@ lerp(const complex<_Tp>& self_val, const complex<_Tp>& end_val, const complex<_T
         ? self_val.real() + weight_val.real() * (end_val.real() - self_val.real())
         : end_val.real() -
             (end_val.real() - self_val.real()) * (static_cast<_Tp>(1) - weight_val.real()))
-	+ ((std::abs(weight_val.imag()) < 0.5)
+        + ((std::abs(weight_val.imag()) < 0.5)
             ? self_val.imag() + weight_val.imag() * (end_val.imag() - self_val.imag())
             : end_val.imag() -
                 (end_val.imag() - self_val.imag()) * (static_cast<_Tp>(1) - weight_val.imag()));
diff --git a/aten/src/ATen/native/cuda/Lerp.cu b/aten/src/ATen/native/cuda/Lerp.cu
@@ -26,6 +26,7 @@ void lerp_tensor_kernel(at::TensorIteratorBase& iter) {
         }
       ); // lerp_tensor_string
   AT_DISPATCH_COMPLEX_TYPES(dtype, "lerp_cuda", [&] {
+        using opmath_t = at::opmath_type<scalar_t>;
         jitted_gpu_kernel<
           /*name=*/ lerp_tensor_name,
           /*return_dtype=*/ scalar_t,
@@ -34,16 +35,20 @@ void lerp_tensor_kernel(at::TensorIteratorBase& iter) {
       });
 #else
   AT_DISPATCH_COMPLEX_TYPES(dtype, "lerp_cuda", [&] {
+      using opmath_t = at::opmath_type<scalar_t>;
       at::native::gpu_kernel(
         iter,
         [] GPU_LAMBDA(
             scalar_t self_val,
             scalar_t end_val,
             scalar_t weight_val) -> scalar_t {
-          return (std:abs(weight_val) < 0.5)
-              ? self_val + weight_val * (end_val - self_val)
-              : end_val -
-                  (end_val - self_val) * (static_cast<scalar_t>(1) - weight_val);
+          opmath_t self_val_f = self_val;
+          opmath_t end_val_f = end_val;
+          opmath_t weight_val_f = weight_val;
+          return (std:abs(weight_val_f) < 0.5)
+              ? self_val_f + weight_val_f * (end_val_f - self_val_f)
+              : end_val_f -
+                  (end_val_f - self_val_f) * (opmath_t{1} - weight_val_f);
         });
       });
 #endif
@@ -96,14 +101,17 @@ void lerp_scalar_kernel(at::TensorIteratorBase& iter, const c10::Scalar& weight)
   });
 #else
   AT_DISPATCH_COMPLEX_TYPES(dtype, "lerp_cuda", [&] {
-      auto weight_val = weight.to<scalar_t>();
+      using opmath_t = at::opmath_type<scalar_t>;
+      auto weight_val = weight.to<opmath_t>();
       gpu_kernel(
         iter,
         [=] GPU_LAMBDA(scalar_t self_val, scalar_t end_val) {
-          return (std::abs(weight_val) < 0.5)
-              ? self_val + weight_val * (end_val - self_val)
-              : end_val -
-                  (end_val - self_val) * (static_cast<scalar_t>(1) - weight_val);
+          opmath_t self_val_f = self_val;
+          opmath_t end_val_f = end_val;
+          return (std::abs(weight_val_f) < 0.5)
+              ? self_val_f + weight_val_f * (end_val_f - self_val_f)
+              : end_val_f -
+                  (end_val_f - self_val_f) * (opmath_t{1} - weight_val_f);
         });
   });
 #endif
