[ROCm] Improvements for vectorized elementwise kernels (pytorch#143269) (#1874)

* Make io_size calculation as minimum of size of input and output size,
rather than the summation of all sizes
* for e.g, for torch.add() on half dtypes (bfloat16/float16),
calc_io_size() returns 6 causing elems_per_thread to be 4
   * But elems_per_thread = 8 works better on half datypes for AMD gpus
* Enable *_load_dwordx4 ISA for 16-bit and 8-bit dtypes on AMD gpus by
using vector size of 8 and 16 respectively

Co-author: @akadutta

Pull Request resolved: pytorch#143269
Approved by: https://github.com/jeffdaily,
https://github.com/pruthvistony

Co-authored-by: Pruthvi Madugundu <[email protected]>
(cherry picked from commit 4686828)

Author: 2 people

aten/src/ATen/native/cuda/CUDAJitLoops.cuh

@@ -49,8 +49,8 @@ struct JittedVecKernelCache {
   at::cuda::jit::NvrtcFunction vec1;
   at::cuda::jit::NvrtcFunction vec2;
   at::cuda::jit::NvrtcFunction vec4;
-  at::cuda::jit::NvrtcFunction vec8;
 #ifdef USE_ROCM
+  at::cuda::jit::NvrtcFunction vec8;
   at::cuda::jit::NvrtcFunction vec16;
 #endif
 
@@ -150,11 +150,11 @@ void launch_jitted_vectorized_kernel(
 #ifdef USE_ROCM
   if (vec_size == 16) {
     fn_ptr = &fn_cache.vec16;
+  } else if (vec_size == 8) {
+    fn_ptr = &fn_cache.vec8;
   } else
 #endif
-  if (vec_size == 8) {
-    fn_ptr = &fn_cache.vec8;
-  } else if (vec_size == 4) {
+  if (vec_size == 4) {
     fn_ptr = &fn_cache.vec4;
   } else if (vec_size == 2) {
     fn_ptr = &fn_cache.vec2;

aten/src/ATen/native/cuda/CUDALoops.cuh

@@ -240,12 +240,12 @@ static inline void launch_vectorized_kernel(
           <<<grid, num_threads(), 0, stream>>>(N, f, data);
       C10_CUDA_KERNEL_LAUNCH_CHECK();
       break;
-#endif
     case 8:
       vectorized_elementwise_kernel<8, func_t, array_t>
           <<<grid, num_threads(), 0, stream>>>(N, f, data);
       C10_CUDA_KERNEL_LAUNCH_CHECK();
       break;
+#endif
     case 4:
       vectorized_elementwise_kernel<4, func_t, array_t>
           <<<grid, num_threads(), 0, stream>>>(N, f, data);

aten/src/ATen/native/cuda/jit_utils.h

@@ -60,7 +60,6 @@ inline int can_vectorize_up_to(size_t default_alignment, void *pointer) {
   if ((default_alignment <= 2) && (ip % (8 * default_alignment) == 0)) {
     return 8;
   }
-#else
   if (ip % (8 * default_alignment) == 0) {
     return 8;
   }