Repository: cupy/cupy

Issue:
Need housekeeping macros for HIP
According to an earlier discussion (https://github.com/cupy/cupy/pull/4367#discussion_r542051342), it's better to collect the macro redefinitions for functions that are not supported by HIP (ex: `__shfl*sync()`) in one place so that they are always correctly included when compiling on HIP. This would benefit CuPy JIT too (ex: #5335). I could try to work on this later next week...


Discussion:

Resolving patch:
diff --git a/cupyx/jit/_builtin_funcs.py b/cupyx/jit/_builtin_funcs.py
--- a/cupyx/jit/_builtin_funcs.py
+++ b/cupyx/jit/_builtin_funcs.py
@@ -126,8 +126,9 @@ def __call__(self, *, mask=0xffffffff):
 
     def call(self, env, *, mask=None):
         if runtime.is_hip:
-            warnings.warn(f'mask {mask} is ignored on HIP', RuntimeWarning)
-            mask = None
+            if mask is not None:
+                warnings.warn(f'mask {mask} is ignored on HIP', RuntimeWarning)
+                mask = None
 
         if mask:
             if isinstance(mask, Constant):


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
@@ -210,8 +210,6 @@ def f(x, y, buf):
         f((1,), (32, 32), (x, y, buf))
         assert bool((x == y).all())
 
-    # TODO(leofang): enable HIP when cupy/cupy#5348 is resolved
-    @unittest.skipIf(runtime.is_hip, 'HIP is not yet supported')
     def test_syncwarp(self):
         @jit.rawkernel()
         def f(x):
@@ -228,8 +226,6 @@ def f(x):
         y[16:] += 1
         assert bool((x == y).all())
 
-    # TODO(leofang): enable HIP when cupy/cupy#5348 is resolved
-    @unittest.skipIf(runtime.is_hip, 'HIP is not yet supported')
     def test_syncwarp_mask(self):
         @jit.rawkernel()
         def f(x, m):
@@ -320,8 +316,6 @@ def f(x, y, size):
         f[5, 6](x, y, numpy.uint32(1024))
         assert bool((x == y).all())
 
-    # TODO(leofang): enable HIP when cupy/cupy#5348 is resolved
-    @unittest.skipIf(runtime.is_hip, 'HIP is not yet supported')
     # TODO(leofang): test float16 ('e') once cupy/cupy#5346 is resolved
     @testing.for_dtypes('iIlqfd' if runtime.is_hip else 'iIlLqQfd')
     def test_shfl(self, dtype):
@@ -340,8 +334,6 @@ def f(a, b):
         f[1, 32](a, b)
         assert (b == a * cupy.ones((32,), dtype=dtype)).all()
 
-    # TODO(leofang): enable HIP when cupy/cupy#5348 is resolved
-    @unittest.skipIf(runtime.is_hip, 'HIP is not yet supported')
     def test_shfl_width(self):
         @jit.rawkernel()
         def f(a, b, w):
@@ -357,8 +349,6 @@ def f(a, b, w):
             c[c % w != 0] = c[c % w == 0]
             assert (b == c).all()
 
-    # TODO(leofang): enable HIP when cupy/cupy#5348 is resolved
-    @unittest.skipIf(runtime.is_hip, 'HIP is not yet supported')
     # TODO(leofang): test float16 ('e') once cupy/cupy#5346 is resolved
     @testing.for_dtypes('iIlqfd' if runtime.is_hip else 'iIlLqQfd')
     def test_shfl_up(self, dtype):
@@ -374,25 +364,24 @@ def f(a):
         expected = [i for i in range(N)] + [i for i in range(32-N)]
         assert(a == cupy.asarray(expected, dtype=dtype)).all()
 
-    # TODO(leofang): enable HIP when cupy/cupy#5348 is resolved
-    @unittest.skipIf(runtime.is_hip, 'HIP is not yet supported')
     # TODO(leofang): test float16 ('e') once cupy/cupy#5346 is resolved
     @testing.for_dtypes('iIlqfd' if runtime.is_hip else 'iIlLqQfd')
     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
 
         @jit.rawkernel()
         def f(a):
             value = jit.shfl_down_sync(0xffffffff, a[jit.threadIdx.x], N)
             a[jit.threadIdx.x] = value
 
-        a = cupy.arange(32, dtype=dtype)
-        f[1, 32](a)
-        expected = [i for i in range(N, 32)] + [(32-N+i) for i in range(N)]
+        a = cupy.arange(block, dtype=dtype)
+        f[1, block](a)
+        expected = [i for i in range(N, block)]
+        expected += [(block-N+i) for i in range(N)]
         assert(a == cupy.asarray(expected, dtype=dtype)).all()
 
-    # TODO(leofang): enable HIP when cupy/cupy#5348 is resolved
-    @unittest.skipIf(runtime.is_hip, 'HIP is not yet supported')
     # TODO(leofang): test float16 ('e') once cupy/cupy#5346 is resolved
     @testing.for_dtypes('iIlqfd' if runtime.is_hip else 'iIlLqQfd')
     def test_shfl_xor(self, dtype):