Amd support

.gitlab-ci.yml

@@ -1,15 +1,20 @@
 variables:
   CUDA_IMAGE: "gitlab-registry.mpcdf.mpg.de/mpcdf/ci-module-image/nvhpcsdk_26:2026"
+  ROCM_IMAGE: "gitlab-registry.mpcdf.mpg.de/mpcdf/ci-module-image/gcc_15-openmpi_5_0-rocm_7_1:2026"
   PIP_DISABLE_PIP_VERSION_CHECK: "1"
   CUNUMPY_BACKEND: "cupy"
   CUNUMPY_REQUIRE_CUDA: "1"
   # empty: the newest CuPy; gpu_tests_min_cupy pins the oldest supported one
   CUPY_PIN: ""
+  # MI210 (MPCDF-GPU-05) is CDNA2, offload architecture gfx90a; also covers
+  # MI200/MI250(X). Change this (and the gpu-amd-gfx* tag below) for another
+  # AMD GPU generation.
+  HCC_AMDGPU_TARGET: "gfx90a"
 
 stages:
   - test
 
-.gpu_job:
+.gpu_job_nvidia:
   stage: test
   image: ${CUDA_IMAGE}
   tags:
@@ -26,25 +31,57 @@ stages:
     - export PATH="$HOME/.local/bin:$PATH"
     - python3 -c "import cupy; cupy.show_config()"
 
-gpu_tests:
-  extends: .gpu_job
+# CuPy has no prebuilt ROCm wheel for recent ROCm (the old cupy-rocm-*
+# packages stop at ROCm 4.3), so this builds it from source against the
+# module-provided ROCm. HCC_AMDGPU_TARGET above must match the runner's GPU.
+.gpu_job_amd:
+  stage: test
+  image: ${ROCM_IMAGE}
+  tags:
+    - gpu-amd
+    - gpu-amd-gfx90a
+  before_script:
+    - module load gcc/15
+    - module load openmpi/5.0
+    - module load rocm/7.1
+    - module load python-waterboa/2025.06
+    - rocm-smi
+    - export ROCM_HOME="${ROCM_PATH:-/opt/rocm}"
+    - export CUPY_INSTALL_USE_HIP=1
+    - python3 -m pip install --user "cupy${CUPY_PIN}"
+    - python3 -m pip install --user -e ".[test]"
+    - export PATH="$HOME/.local/bin:$PATH"
+    - python3 -c "import cupy; cupy.show_config()"
+    - python3 -c "import cupy; assert cupy.cuda.runtime.is_hip, 'expected a HIP build of CuPy'"
+
+gpu_tests_nvidia:
+  extends: .gpu_job_nvidia
+  script:
+    - python3 -m pytest -q tests
+    - python3 testing/cupyx_names.py
+
+# cunumpy/reduce.cuh and cunumpy/scan.cuh are not yet ported to HIP: tests that
+# need them use `requires_warp_shuffle`, which skips (not fails) here even
+# though CUNUMPY_REQUIRE_CUDA is set, see cunumpy.kernel_testing.
+gpu_tests_amd:
+  extends: .gpu_job_amd
   script:
     - python3 -m pytest -q tests
     - python3 testing/cupyx_names.py
 
 # The oldest supported CuPy (cunumpy.MIN_CUPY_VERSION): the whole suite, and the
 # fake cupyx.scipy must list exactly its names (_fake_cupyx.NAMES).
-gpu_tests_min_cupy:
-  extends: .gpu_job
+gpu_tests_min_cupy_nvidia:
+  extends: .gpu_job_nvidia
   variables:
     CUPY_PIN: "==14.0.0"
   script:
     - python3 -c "import cupy, cunumpy; assert cupy.__version__ == cunumpy.MIN_CUPY_VERSION, cupy.__version__"
     - python3 -m pytest -q tests
     - python3 testing/cupyx_names.py --exact
 
-gpu_sanitizers:
-  extends: .gpu_job
+gpu_sanitizers_nvidia:
+  extends: .gpu_job_nvidia
   timeout: 45 minutes
   parallel:
     matrix:

docs/source/kernels/cuda-headers.md

@@ -19,6 +19,30 @@ in [Writing CUDA kernels](cuda-kernel.md)).
 | `cunumpy/random.cuh` | counter-based random numbers, equal to `xp.rng.philox_*` |
 | `cunumpy/morton.cuh` | Morton (Z-order) keys, equal to `xp.algorithms.morton_*` |
 
+**AMD/ROCm.** CuPy's own API is the same on a ROCm build, so `xp.set_backend("cupy")`
+and the device/memory/stream helpers of `xp.cuda` work unchanged; check
+`xp.cuda.is_hip()` to tell the two apart where it matters. `CudaKernel`
+compiles plain CUDA C through `cupy.RawKernel`, which HIPRTC accepts for
+most kernels. Three known gaps, found on real HIP/ROCm hardware (an AMD
+Instinct MI210) and all raising `NotImplementedError` at
+`CudaKernel.compile()` on HIP rather than silently misbehaving, none yet
+root-caused:
+
+* `cunumpy/reduce.cuh` (and `cunumpy/scan.cuh`, which includes it), see below.
+* A kernel taking an `Array5D`/`CArray5D` view or higher (`Array1D` to
+  `Array4D` are unaffected): launching reliably corrupts the device, even for
+  an in-bounds access (a HIP code-generation or by-value struct-argument
+  issue in the generic, variadic-template `ArrayView<T, N>` is suspected).
+* A kernel that includes `cunumpy/array_view.cuh` and compiles with
+  `CUNUMPY_BOUNDS_CHECK` (directly, or via debug mode, which adds it): merely
+  compiling the resulting `printf()` and trap reliably corrupts the device on
+  any launch, even one that never takes an out-of-bounds index (a HIP
+  device-`printf` or trap-instruction issue is suspected).
+
+Use `cunumpy.kernel_testing`'s `requires_warp_shuffle`, `requires_high_dim_cuda_views`
+and `requires_bounds_check_views` markers (or their `*_available()` functions)
+to skip tests that hit these on HIP, the way cunumpy's own test suite does.
+
 ## `cunumpy/index.cuh`
 
 ```c
@@ -113,6 +137,14 @@ gets the result.
   default requires all 32 lanes. These rules follow [NVIDIA's warp intrinsic
   constraints](https://docs.nvidia.com/cuda/cuda-programming-guide/05-appendices/cpp-language-extensions.html).
 
+**Not ported to HIP/ROCm.** This header hardcodes a 32-lane warp and CUDA's
+`_sync` shuffle intrinsics; AMD wavefronts are commonly 64 lanes wide (CDNA:
+MI100/MI200/MI300), and HIP's shuffles have no mask argument. Rather than
+compile silently-wrong reductions, `CudaKernel.compile()` raises
+`NotImplementedError` for a kernel that includes `cunumpy/reduce.cuh` (or
+`cunumpy/scan.cuh`, which includes it) when the active CuPy build targets
+HIP (`xp.cuda.is_hip()` is True).
+
 ```c
 extern "C" __global__ void kinetic_energy(const double* v, long long n,
                                           double mass, double* energy) {

src/cunumpy/_cuda_kernel.py

@@ -857,6 +857,17 @@ def check(value: Any) -> Any:
     return check
 
 
+def _max_view_ndim(params: Iterable[CudaParameter]) -> int:
+    """Return the largest `view_ndim` among `params` and their struct fields (0 if none)."""
+    best = 0
+    for p in params:
+        if p.view_ndim is not None:
+            best = max(best, p.view_ndim)
+        if p.struct is not None:
+            best = max(best, _max_view_ndim(p.struct.fields))
+    return best
+
+
 def _checker(param: CudaParameter, index: int) -> Callable[[Any], Any]:
     if param.struct is not None:
         return _struct_checker(param, index)
@@ -2183,13 +2194,15 @@ def included_headers(self) -> tuple[Path, ...]:
         )
 
     def compile_options(self) -> tuple[str, ...]:
-        """Return the NVRTC options a compilation now would use.
+        """Return the compiler options a compilation now would use.
 
         These are :attr:`options`, ``-I`` for cunumpy's header directory,
-        :data:`~cunumpy.cuda.DEBUG_OPTIONS` if :meth:`debug_active`, and, if
-        the source includes headers, ``-DCUNUMPY_INCLUDE_HASH=0x<hash>`` of
-        their contents (:func:`~cunumpy.cuda.include_hash`). CuPy keys its
-        kernel cache on the options, so a changed header means a recompile.
+        :data:`~cunumpy.cuda.DEBUG_OPTIONS` if :meth:`debug_active` (minus
+        ``-lineinfo`` on a HIP/ROCm build, an NVRTC-only flag HIPRTC does not
+        accept), and, if the source includes headers,
+        ``-DCUNUMPY_INCLUDE_HASH=0x<hash>`` of their contents
+        (:func:`~cunumpy.cuda.include_hash`). CuPy keys its kernel cache on
+        the options, so a changed header means a recompile.
 
         Returns
         -------
@@ -2202,7 +2215,12 @@ def compile_options(self) -> tuple[str, ...]:
         if cunumpy_include not in options:
             options += (cunumpy_include,)
         if self.debug_active():
-            options += tuple(o for o in DEBUG_OPTIONS if o not in options)
+            from cunumpy._device import is_hip
+
+            debug_options = DEBUG_OPTIONS
+            if is_hip():
+                debug_options = tuple(o for o in debug_options if o != "-lineinfo")
+            options += tuple(o for o in debug_options if o not in options)
         headers = self.included_headers
         if headers:
             options += (f"-DCUNUMPY_INCLUDE_HASH=0x{include_hash(headers)}",)
@@ -2287,6 +2305,27 @@ def compile(self, *, log_stream: Any = None) -> Any:
         ------
         RuntimeError
             If CuPy or a GPU is not available.
+        NotImplementedError
+            If the source includes ``cunumpy/reduce.cuh`` or
+            ``cunumpy/scan.cuh`` and the active CuPy build targets HIP/ROCm:
+            their warp-level shuffles assume a 32-lane warp and CUDA's
+            ``_sync`` intrinsics, neither of which hold on AMD wavefronts
+            (commonly 64 lanes on CDNA GPUs), so compiling them as-is would
+            silently reduce or scan the wrong set of threads. Also raised, on
+            a HIP/ROCm build, for a kernel using an ``Array5D``/``CArray5D``
+            view or higher: launching one reliably corrupts the device there
+            (observed as ``hipErrorIllegalState`` on a later launch), even for
+            an in-bounds access; ``Array1D`` to ``Array4D`` are unaffected.
+            Not yet root-caused (looks like a HIP code-generation or
+            by-value-struct-argument issue in the generic, variadic-template
+            ``ArrayView<T, N>`` that implements dimensions 5 and up, see
+            ``cunumpy/array_view.cuh``). Also raised, on a HIP/ROCm build,
+            for a kernel that includes ``cunumpy/array_view.cuh`` and
+            compiles with ``-DCUNUMPY_BOUNDS_CHECK`` (directly or via debug
+            mode): merely compiling the resulting ``printf()`` and trap
+            reliably corrupts the device on any launch of that kernel, even
+            one that never takes an out-of-bounds index; also not yet
+            root-caused.
         """
         from cunumpy.xp import cupy_available
 
@@ -2297,9 +2336,56 @@ def compile(self, *, log_stream: Any = None) -> Any:
             )
         import cupy as cp
 
+        from cunumpy._device import is_hip
+
         device = _current_device()
         if device not in self._compiled:
+            if is_hip():
+                unported = {"reduce.cuh", "scan.cuh"} & {
+                    h.name for h in self.included_headers
+                }
+                if unported:
+                    names = ", ".join(sorted(f"cunumpy/{n}" for n in unported))
+                    raise NotImplementedError(
+                        f"CUDA kernel {self.expression!r} includes {names}, which are "
+                        "not yet ported to HIP/ROCm: their warp-level shuffles assume "
+                        "a 32-lane warp and CUDA's `_sync` intrinsics, neither of "
+                        "which hold on AMD wavefronts (commonly 64 lanes on CDNA "
+                        "GPUs). Compiling them as-is would silently reduce/scan the "
+                        "wrong set of threads rather than fail loudly.",
+                    )
+                if self._signature is not None:
+                    max_view = _max_view_ndim(self._signature)
+                    if max_view >= 5:
+                        raise NotImplementedError(
+                            f"CUDA kernel {self.expression!r} uses an "
+                            f"Array{max_view}D/CArray{max_view}D view: on this "
+                            "HIP/ROCm build, launching a kernel with a 5+ "
+                            "dimensional array view reliably corrupts the device "
+                            "(hipErrorIllegalState on a later launch), even for an "
+                            "in-bounds access. Array1D to Array4D are unaffected; "
+                            "this is not yet root-caused (likely a HIP code "
+                            "generation or by-value struct argument issue in the "
+                            "generic ArrayView<T, N> template, not something "
+                            "cunumpy can safely paper over).",
+                        )
             options = self.compile_options()
+            if (
+                is_hip()
+                and "-DCUNUMPY_BOUNDS_CHECK" in options
+                and "array_view.cuh" in {h.name for h in self.included_headers}
+            ):
+                raise NotImplementedError(
+                    f"CUDA kernel {self.expression!r} includes cunumpy/array_view.cuh "
+                    "and compiles with -DCUNUMPY_BOUNDS_CHECK (directly, or via debug "
+                    "mode) on this HIP/ROCm build: merely compiling the resulting "
+                    "printf() + trap in its bounds check reliably corrupts the "
+                    "device on any launch of that kernel, even one that never takes "
+                    "an out-of-bounds index. Not yet root-caused (a HIP "
+                    "device-printf or trap-instruction issue is suspected); compile "
+                    "without CUNUMPY_BOUNDS_CHECK (and without debug mode, which adds "
+                    "it) for a kernel using array views on this build.",
+                )
             if self._template_args is None:
                 raw = cp.RawKernel(
                     self._source,

src/cunumpy/_device.py

@@ -64,6 +64,38 @@ def device_count() -> int:
         return 0
 
 
+def is_hip() -> bool:
+    """Return whether the active CuPy build targets AMD ROCm/HIP.
+
+    CuPy's own API (``cp.cuda.Device``, ``memory_info``, streams, ...) is the
+    same on a ROCm build; this only distinguishes the GPU vendor for the
+    parts that differ, such as :meth:`~cunumpy.kernels.CudaKernel.compile_options`
+    (NVRTC-only flags) and the warp-level primitives of
+    ``cunumpy/reduce.cuh`` and ``cunumpy/scan.cuh`` (not yet ported to HIP,
+    see :doc:`/kernels/cuda-headers`).
+
+    Returns
+    -------
+    bool
+        True if CuPy is available and reports a HIP runtime, False for a
+        CUDA build or when CuPy/a GPU is unavailable.
+
+    Examples
+    --------
+    >>> xp.cuda.is_hip()  # on a CUDA build, or without CuPy
+    False
+    """
+    if not cupy_available():
+        return False
+
+    import cupy as cp
+
+    try:
+        return bool(cp.cuda.runtime.is_hip)
+    except Exception:  # noqa: BLE001 - tolerate any driver/runtime failure
+        return False
+
+
 def set_device_for_rank(rank: int, devices_per_node: int | None = None) -> int:
     """Select device ``rank % devices_per_node`` for an MPI rank.
 
@@ -200,9 +232,21 @@ def max_shared_memory_per_block(
         return DEFAULT_SHARED_MEMORY_PER_BLOCK
     import cupy as cp
 
-    dev = cp.cuda.Device() if device is None else cp.cuda.Device(device)
-    key = "MaxSharedMemoryPerBlockOptin" if opt_in else "MaxSharedMemoryPerBlock"
-    return int(dev.attributes.get(key, DEFAULT_SHARED_MEMORY_PER_BLOCK))
+    # `cp.cuda.Device().attributes` is unreliable here on a HIP/ROCm build (it
+    # can report 0 for these keys); `getDeviceProperties()`'s `sharedMemPerBlock`
+    # and `sharedMemPerBlockOptin` fields are the ones `CudaKernel` itself
+    # already relies on for the real launch-time check, on both backends.
+    dev_id = cp.cuda.Device().id if device is None else device
+    properties = cp.cuda.runtime.getDeviceProperties(dev_id)
+    base = (
+        int(properties.get("sharedMemPerBlock", 0)) or DEFAULT_SHARED_MEMORY_PER_BLOCK
+    )
+    if not opt_in:
+        return base
+    # The "opt in" larger dynamic-shared-memory limit is an NVIDIA concept
+    # (cudaFuncAttributeMaxDynamicSharedMemorySize); on a HIP build there is no
+    # larger limit to opt into, so the base limit is also the opt-in one.
+    return max(base, int(properties.get("sharedMemPerBlockOptin", 0)))
 
 
 def free_memory() -> None:

src/cunumpy/cuda/__init__.py

@@ -35,6 +35,7 @@
     device_count,
     free_memory,
     get_cuda_debug,
+    is_hip,
     max_shared_memory_per_block,
     memory_info,
     pin_memory,
@@ -69,6 +70,7 @@
     "free_memory",
     "get_cuda_debug",
     "include_hash",
+    "is_hip",
     "max_shared_memory_per_block",
     "memory_info",
     "parse_cuda_signature",

src/cunumpy/cuda/include/cunumpy/array_view.cuh

@@ -42,13 +42,23 @@
 #define CUNUMPY_ARRAY_VIEW_CUH
 
 #ifdef CUNUMPY_BOUNDS_CHECK
+// __trap() is a CUDA device builtin, found unchanged by NVRTC and (deliberately
+// shadowable) by the C++ compiler used for CPU emulation (cunumpy._emulation
+// defines its own host __trap()). HIPRTC has no such builtin; only a real HIP
+// device compile (never the CPU emulation path, which never defines these
+// macros) uses __builtin_trap() instead.
+#if defined(__HIP_DEVICE_COMPILE__) || defined(__HIP_PLATFORM_AMD__) || defined(__HIPCC__)
+#define CUNUMPY_TRAP() __builtin_trap()
+#else
+#define CUNUMPY_TRAP() __trap()
+#endif
 #define CUNUMPY_CHECK_INDEX(index, axis, extent)                                 \
     do {                                                                         \
         if ((index) < 0 || (index) >= (extent)) {                                \
             printf("cunumpy: index %lld is out of bounds for axis %d with size " \
                    "%lld\n",                                                     \
                    (long long)(index), (int)(axis), (long long)(extent));        \
-            __trap();                                                            \
+            CUNUMPY_TRAP();                                                      \
         }                                                                        \
     } while (0)
 #else

src/cunumpy/cuda/include/cunumpy/reduce.cuh

@@ -27,6 +27,12 @@
 //
 // The warp size is 32, as on all NVIDIA GPUs.
 //
+// Not yet ported to HIP/ROCm: AMD wavefronts are commonly 64 lanes (CDNA:
+// MI100/MI200/MI300) rather than 32, and HIP's shuffle intrinsics have no
+// `_sync`/mask variant. `CudaKernel.compile()` raises `NotImplementedError`
+// for a kernel that includes this header on a HIP build, instead of silently
+// reducing/scanning the wrong set of threads.
+//
 // The header directory is added to every CudaKernel's NVRTC options.
 
 #ifndef CUNUMPY_REDUCE_CUH

src/cunumpy/cuda/include/cunumpy/scan.cuh

@@ -4,6 +4,7 @@
 // every thread must call a block function, including padding threads (value 0).
 // Supports CUDA shuffle arithmetic types, arbitrary masks and partial warps.
 // Block scans use 32 shared values per type; calls must be collective.
+// Not yet ported to HIP/ROCm, see cunumpy/reduce.cuh.
 #ifndef CUNUMPY_SCAN_CUH
 #define CUNUMPY_SCAN_CUH