Use explicit templates in gpu_kernel_with_scalars (pytorch#40992)

malfet · csarofeen · commit ead6fe4b448f · 2020-07-07T09:07:27.000-04:00
Summary: This trick should have no effect on performance, but it reduces size of kernels using the template by 10% For example, sizeof(BinaryMulDivKernel.cu.o) compiled by CUDA-10.1 toolchain for sm_75 before the change was 4.2Mb, after 3.8Mb Pull Request resolved: pytorch#40992 Differential Revision: D22398733 Pulled By: malfet fbshipit-source-id: 6576f4da00dc5fc2575b2313577f52c6571d5e6f
diff --git a/aten/src/ATen/native/cuda/Loops.cuh b/aten/src/ATen/native/cuda/Loops.cuh
@@ -55,7 +55,6 @@ namespace at { namespace native {
 
 template <typename func_t>
 void gpu_kernel(TensorIterator& iter, const func_t& f) {
-  ASSERT_HOST_DEVICE_LAMBDA(func_t);
 
   for (int arg = 0; arg < iter.ntensors(); arg++) {
     TORCH_INTERNAL_ASSERT(iter.device(arg).is_cuda());
@@ -75,6 +74,36 @@ void gpu_kernel(TensorIterator& iter, const func_t& f) {
   gpu_kernel_impl(iter, f);
 }
 
+template<typename func_t>
+struct AUnaryFunctor {
+  using traits = function_traits<func_t>;
+  using arg1_t = typename traits::template arg<0>::type;
+  using arg2_t = typename traits::template arg<1>::type;
+  using return_t = typename traits::result_type;
+  __device__ return_t operator()(arg2_t b) const {
+    return f(a, b);
+  }
+  AUnaryFunctor(func_t f_, arg1_t a_): f(f_), a(a_) {}
+  private:
+    func_t f;
+    arg1_t a;
+};
+
+template<typename func_t>
+struct BUnaryFunctor {
+  using traits = function_traits<func_t>;
+  using arg1_t = typename traits::template arg<0>::type;
+  using arg2_t = typename traits::template arg<1>::type;
+  using return_t = typename traits::result_type;
+  __device__ return_t operator()(arg1_t a) const {
+    return f(a, b);
+  }
+  BUnaryFunctor(func_t f_, arg2_t b_): f(f_), b(b_) {}
+  private:
+    func_t f;
+    arg2_t b;
+};
+
 template <typename func_t>
 void gpu_kernel_with_scalars(TensorIterator& iter, const func_t& f) {
   ASSERT_HOST_DEVICE_LAMBDA(func_t);
@@ -85,23 +114,17 @@ void gpu_kernel_with_scalars(TensorIterator& iter, const func_t& f) {
       traits::arity == 2,
       "gpu_kernel_with_scalars only supports two input arguments");
 
+  using arg1_t = typename traits::template arg<0>::type;
+  using arg2_t = typename traits::template arg<1>::type;
   if (iter.is_cpu_scalar(1)) {
-    using arg1_t = typename traits::template arg<0>::type;
-    using arg2_t = typename traits::template arg<1>::type;
-    auto a = iter.scalar_value<arg1_t>(1);
+    AUnaryFunctor<func_t> af(f, iter.scalar_value<arg1_t>(1));
     iter.remove_operand(1);
     const OptionalDeviceGuard device_guard(device_of(iter.tensor(1)));
-    gpu_kernel(iter, [=]GPU_LAMBDA(arg2_t b) {
-      return f(a, b);
-    });
+    gpu_kernel(iter, af);
   } else if (iter.is_cpu_scalar(2)) {
-    using arg1_t = typename traits::template arg<0>::type;
-    using arg2_t = typename traits::template arg<1>::type;
-    auto b = iter.scalar_value<arg2_t>(2);
+    BUnaryFunctor<func_t> bf(f, iter.scalar_value<arg2_t>(2));
     iter.remove_operand(2);
-    gpu_kernel(iter, [=]GPU_LAMBDA(arg1_t a) {
-      return f(a, b);
-    });
+    gpu_kernel(iter, bf);
   } else {
     gpu_kernel(iter, f);
   }
