Repository: cupy/cupy

Issue:
Use `warpSize` to get the warp size in ROCm environment
CuPy, currently, has a hard-coded warp size of 64 for ROCm environment. However, AMD gpu's warp size is not necessarily 64; for example, Navi21 that ROCm 5.0 supports has the warp size of 32. The following can be a good solution.

https://github.com/pytorch/pytorch/blob/2cefbb71cf65696e28ee1bdfce06d6846285e5fb/c10/macros/Macros.h#L322

Thanks @amathews-amd for pointing this out.


Discussion:
<!-- Added by Auto Notify Bot (hip) -->
cc/ @amathews-amd
Are there any code in CuPy that assumes warpSize = 64?

edit: @asi1024 found this one :smile:
https://github.com/cupy/cupy/blob/eabb8f9792310522f11a6a67f9aab33249f64047/cupyx/jit/_interface.py#L186-L187
Also,
https://github.com/cupy/cupy/blob/eabb8f9792310522f11a6a67f9aab33249f64047/cupy/_core/_routines_math.pyx#L456
Bumped the priority to high (https://github.com/cupy/cupy/issues/7234#issuecomment-1377921410). Maybe we can just replace it with `warpSize` for both CUDA and ROCm? I guess PyTorch hard-codes it for compile-time optimization but I think CuPy don't have to care about it as we dynamically compile code into PTX/CUBIN and host compilers are not involved.
@kmaehashi 
Here is my naive solution that worked for what I needed, https://github.com/cupy/cupy/commit/eda2b396a5255e745fc8a5caaa0d7cc621f16073

Let me know if you need anything else from me to get a working solution implemented.
Thanks @ajflood! As for `_routines_math.pyx`, `cupy.cuda.device._get_device().attributes` could be better as it caches. JIT is the difficult part as CUDA API calls should not be happen when `import cupy`. Maybe @asi1024 can better follow up here.

Resolving patch:
diff --git a/cupy/_core/__init__.py b/cupy/_core/__init__.py
--- a/cupy/_core/__init__.py
+++ b/cupy/_core/__init__.py
@@ -18,6 +18,7 @@
 from cupy._core._kernel import create_ufunc  # NOQA
 from cupy._core._kernel import ElementwiseKernel  # NOQA
 from cupy._core._kernel import ufunc  # NOQA
+from cupy._core._kernel import _get_warpsize  # NOQA
 from cupy._core._reduction import create_reduction_func  # NOQA
 from cupy._core._reduction import ReductionKernel  # NOQA
 from cupy._core._routines_binary import bitwise_and  # NOQA
diff --git a/cupyx/jit/_interface.py b/cupyx/jit/_interface.py
--- a/cupyx/jit/_interface.py
+++ b/cupyx/jit/_interface.py
@@ -3,7 +3,6 @@
 
 import numpy
 
-from cupy_backends.cuda.api import runtime
 import cupy
 from cupy._core import core
 from cupyx.jit import _compile
@@ -182,11 +181,8 @@ def wrapper(func):
 blockIdx = _internal_types.Data('blockIdx', _cuda_types.dim3)
 gridDim = _internal_types.Data('gridDim', _cuda_types.dim3)
 
-warpsize = _internal_types.Data(
-    '64' if runtime.is_hip else '32', _cuda_types.uint32)
+warpsize = _internal_types.Data('warpSize', _cuda_types.int32)
 warpsize.__doc__ = r"""Returns the number of threads in a warp.
 
-In CUDA this is always 32, and in ROCm/HIP always 64.
-
 .. seealso:: :obj:`numba.cuda.warpsize`
 """


Test patch:
diff --git a/tests/cupyx_tests/jit_tests/test_raw.py b/tests/cupyx_tests/jit_tests/test_raw.py
--- a/tests/cupyx_tests/jit_tests/test_raw.py
+++ b/tests/cupyx_tests/jit_tests/test_raw.py
@@ -580,7 +580,7 @@ def f(a):
     def test_shfl_down(self, dtype):
         N = 5
         # __shfl_down() on HIP does not seem to have the same behavior...
-        block = 64 if runtime.is_hip else 32
+        block = cupy._core._get_warpsize()
 
         @jit.rawkernel()
         def f(a):
@@ -627,7 +627,7 @@ def f(arr):
             if x < arr.size:
                 arr[x] = jit.laneid()
 
-        N = 64 if runtime.is_hip else 32
+        N = cupy._core._get_warpsize()
         x = cupy.zeros((N*2,), dtype=cupy.uint32)
         f((1,), (N*2,), (x,))
         y = cupy.arange(N*2, dtype=cupy.uint32) % N
@@ -640,7 +640,7 @@ def f(arr):
             if x == 0:
                 arr[0] = jit.warpsize
 
-        N = 64 if runtime.is_hip else 32
+        N = cupy._core._get_warpsize()
         x = cupy.zeros((1,), dtype=cupy.uint32)
         f((1,), (1,), (x,))
         assert x == N