cherry-pick fix for ROCm 6.4

.github/workflows/pretest.yml

@@ -30,6 +30,7 @@ jobs:
 
     - name: Setup pre-commit
       run: |
+        pip install --upgrade autopep8
         pip install pre-commit
 
     - name: Check
@@ -106,4 +107,4 @@ jobs:
 
     - name: Build & Test
       run: |
-        docker run --rm -v "${PWD}:/src" -w /src "rocm/dev-ubuntu-22.04:6.0" bash .github/workflows/scripts/pretest-rocm-test.sh
+        docker run --rm -v "${PWD}:/src" -w /src "rocm/dev-ubuntu-22.04:6.4" bash .github/workflows/scripts/pretest-rocm-test.sh

cupy/_core/core.pyx

@@ -261,9 +261,6 @@ cdef class _ndarray_base:
 
     @property
     def __cuda_array_interface__(self):
-        if runtime._is_hip_environment:
-            raise AttributeError(
-                'HIP/ROCm does not support cuda array interface')
         cdef dict desc = {
             'shape': self.shape,
             'typestr': self.dtype.str,
@@ -1800,8 +1797,6 @@ cdef class _ndarray_base:
                 # `_kernel._preprocess_args`.
                 check = (hasattr(x, '__cuda_array_interface__')
                          or hasattr(x, '__cupy_get_ndarray__'))
-                if runtime._is_hip_environment and isinstance(x, ndarray):
-                    check = True
                 if (not check
                         and not type(x) in _scalar.scalar_type_set
                         and not isinstance(x, numpy.ndarray)):
@@ -3000,9 +2995,6 @@ cpdef _ndarray_base asfortranarray(_ndarray_base a, dtype=None):
 
 
 cpdef _ndarray_base _convert_object_with_cuda_array_interface(a):
-    if runtime._is_hip_environment:
-        raise RuntimeError(
-            'HIP/ROCm does not support cuda array interface')
 
     cdef Py_ssize_t sh, st
     cdef dict desc = a.__cuda_array_interface__

cupy/cuda/cupy_cufft.h

@@ -12,7 +12,12 @@
 #include <cufftXt.h>
 
 #elif defined(CUPY_USE_HIP)
+#include <hip/hip_version.h> //for HIP_VERSION
+#if HIP_VERSION >= 50530600 
+#include <hipfft/hipfft.h>
+#else
 #include <hipfft.h>
+#endif
 
 extern "C" {
 

cupy/cuda/cupy_thrust.cu

@@ -476,7 +476,7 @@ struct _argsort {
                       #else
                       thrust::make_counting_iterator<size_t>(0),
                       thrust::make_counting_iterator<size_t>(size),
-                      thrust::make_constant_iterator<ptrdiff_t>(shape[ndim-1]),
+		      thrust::make_constant_iterator<ptrdiff_t>(shape[ndim-1]),
                       #endif
                       dp_keys_first,
                       thrust::divides<size_t>());

cupy/cuda/cupy_thrust.h

@@ -4,7 +4,14 @@
 #ifndef CUPY_NO_CUDA
 #include <vector>
 #include <cstdint>
+
+#ifndef CUPY_USE_HIP
 #include <thrust/version.h>  // for THRUST_VERSION
+#else
+// WAR #9098:
+// rocThrust 3.3.0 (ROCm 6.4.0) cannot be compiled by host compiler
+#define THRUST_VERSION 0
+#endif
 
 void thrust_sort(int, void *, size_t *, const std::vector<ptrdiff_t>&, intptr_t, void *);
 void thrust_lexsort(int, size_t *, void *, size_t, size_t, intptr_t, void *);

cupy/cuda/memory.pyx

@@ -506,10 +506,7 @@ cdef class MemoryPointer:
 
         """
         stream_ptr = stream_module.get_current_stream_ptr()
-        if (
-            not runtime._is_hip_environment
-            and runtime.streamIsCapturing(stream_ptr)
-        ):
+        if runtime.streamIsCapturing(stream_ptr):
             raise RuntimeError(
                 'the current stream is capturing, so synchronous API calls '
                 'are disallowed')
@@ -553,10 +550,7 @@ cdef class MemoryPointer:
 
         """
         stream_ptr = stream_module.get_current_stream_ptr()
-        if (
-            not runtime._is_hip_environment
-            and runtime.streamIsCapturing(stream_ptr)
-        ):
+        if runtime.streamIsCapturing(stream_ptr):
             raise RuntimeError(
                 'the current stream is capturing, so synchronous API calls '
                 'are disallowed')
@@ -580,10 +574,7 @@ cdef class MemoryPointer:
             stream_ptr = stream_module.get_current_stream_ptr()
         else:
             stream_ptr = stream.ptr
-        if (
-            not runtime._is_hip_environment
-            and runtime.streamIsCapturing(stream_ptr)
-        ):
+        if runtime.streamIsCapturing(stream_ptr):
             raise RuntimeError(
                 'the current stream is capturing, so H2D transfers '
                 'are disallowed')
@@ -651,10 +642,7 @@ cdef class MemoryPointer:
 
         """
         stream_ptr = stream_module.get_current_stream_ptr()
-        if (
-            not runtime._is_hip_environment
-            and runtime.streamIsCapturing(stream_ptr)
-        ):
+        if runtime.streamIsCapturing(stream_ptr):
             raise RuntimeError(
                 'the current stream is capturing, so synchronous API calls '
                 'are disallowed')
@@ -678,10 +666,7 @@ cdef class MemoryPointer:
             stream_ptr = stream_module.get_current_stream_ptr()
         else:
             stream_ptr = stream.ptr
-        if (
-            not runtime._is_hip_environment
-            and runtime.streamIsCapturing(stream_ptr)
-        ):
+        if runtime.streamIsCapturing(stream_ptr):
             raise RuntimeError(
                 'the current stream is capturing, so D2H transfers '
                 'are disallowed')
@@ -705,10 +690,7 @@ cdef class MemoryPointer:
 
         """
         stream_ptr = stream_module.get_current_stream_ptr()
-        if (
-            not runtime._is_hip_environment
-            and runtime.streamIsCapturing(stream_ptr)
-        ):
+        if runtime.streamIsCapturing(stream_ptr):
             raise RuntimeError(
                 'the current stream is capturing, so synchronous API calls '
                 'are disallowed')

cupy/cuda/stream.pyx

@@ -369,8 +369,6 @@ class _BaseStream:
             https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__STREAM.html#group__CUDART__STREAM_1g793d7d4e474388ddfda531603dc34aa3
 
         """
-        if runtime._is_hip_environment:
-            raise RuntimeError('This function is not supported on HIP')
         if self.ptr == 0 or self.ptr == 1:
             raise RuntimeError('cannot capture on the default (legacy) stream')
         if mode is None:
@@ -399,8 +397,6 @@ class _BaseStream:
             https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__STREAM.html#group__CUDART__STREAM_1gf5a0efebc818054ceecd1e3e5e76d93e
 
         """
-        if runtime._is_hip_environment:
-            raise RuntimeError('This function is not supported on HIP')
         cdef intptr_t g = runtime.streamEndCapture(self.ptr)
         return graph.Graph.from_stream(g)
 
@@ -415,9 +411,6 @@ class _BaseStream:
                 Programming Guide for detail.
 
         """
-        # TODO(leofang): is it better to be a property?
-        if runtime._is_hip_environment:
-            raise RuntimeError('This function is not supported on HIP')
         try:
             return runtime.streamIsCapturing(self.ptr)
         except RuntimeError:  # can be RuntimeError or CUDARuntimeError
@@ -483,9 +476,6 @@ class Stream(_BaseStream):
             ptr = 0
             device_id = -1
         elif ptds:
-            if runtime._is_hip_environment:
-                raise ValueError('HIP does not support per-thread '
-                                 'default stream (ptds)')
             ptr = runtime.streamPerThread
             device_id = -1
         else:
@@ -544,5 +534,4 @@ class ExternalStream(_BaseStream):
 
 
 Stream.null = Stream(null=True)
-if not runtime._is_hip_environment:
-    Stream.ptds = Stream(ptds=True)
+Stream.ptds = Stream(ptds=True)

cupy/random/cupy_distributions.cuh

@@ -34,7 +34,9 @@ struct rk_binomial_state {
 // When compiling cython extensions with hip 4.0
 // gcc will be used, but the hiprand_kernel can only be compiled with llvm
 // so we need to explicitly declare stubs for the functions
-#if HIP_VERSION > 400
+#if HIP_VERSION >= 50530600
+#include <hiprand/hiprand_kernel.h>
+#elif HIP_VERSION > 400
 #include <hiprand_kernel.h>
 #else
 #include <hiprand.h>

cupy_backends/cuda/api/runtime.pyx

@@ -530,8 +530,8 @@ cpdef intptr_t mallocArray(intptr_t descPtr, size_t width, size_t height,
 
 cpdef intptr_t mallocAsync(size_t size, intptr_t stream) except? 0:
     cdef void* ptr
-    if _is_hip_environment:
-        raise RuntimeError('HIP does not support mallocAsync')
+    if _is_hip_environment and 0 < CUPY_HIP_VERSION < 60200000:
+        raise RuntimeError('mallocAsync is supported since ROCm 6.2')
     with nogil:
         status = cudaMallocAsync(&ptr, size, <driver.Stream>stream)
     check_status(status)
@@ -581,8 +581,8 @@ cpdef freeArray(intptr_t ptr):
     check_status(status)
 
 cpdef freeAsync(intptr_t ptr, intptr_t stream):
-    if _is_hip_environment:
-        raise RuntimeError('HIP does not support freeAsync')
+    if _is_hip_environment and 0 < CUPY_HIP_VERSION < 60200000:
+        raise RuntimeError('freeAsync is supported since ROCm 6.2')
     with nogil:
         status = cudaFreeAsync(<void*>ptr, <driver.Stream>stream)
     check_status(status)
@@ -715,7 +715,7 @@ cpdef PointerAttributes pointerGetAttributes(intptr_t ptr):
     cdef _PointerAttributes attrs
     status = cudaPointerGetAttributes(&attrs, <void*>ptr)
     check_status(status)
-    IF CUPY_CUDA_VERSION > 0:
+    IF CUPY_CUDA_VERSION > 0 or CUPY_HIP_VERSION > 60000000:
         return PointerAttributes(
             attrs.device,
             <intptr_t>attrs.devicePointer,
@@ -896,9 +896,6 @@ cdef _HostFnFunc(void* func_arg) with gil:
 
 cpdef streamAddCallback(intptr_t stream, callback, intptr_t arg,
                         unsigned int flags=0):
-    if _is_hip_environment and stream == 0:
-        raise RuntimeError('HIP does not allow adding callbacks to the '
-                           'default (null) stream')
     func_arg = (callback, arg)
     cpython.Py_INCREF(func_arg)
     with nogil:
@@ -909,9 +906,6 @@ cpdef streamAddCallback(intptr_t stream, callback, intptr_t arg,
 
 
 cpdef launchHostFunc(intptr_t stream, callback, intptr_t arg):
-    if _is_hip_environment:
-        raise RuntimeError('This feature is not supported on HIP')
-
     func_arg = (callback, arg)
     cpython.Py_INCREF(func_arg)
     with nogil:
@@ -933,8 +927,6 @@ cpdef streamWaitEvent(intptr_t stream, intptr_t event, unsigned int flags=0):
 
 
 cpdef streamBeginCapture(intptr_t stream, int mode=streamCaptureModeRelaxed):
-    if _is_hip_environment:
-        raise RuntimeError('streamBeginCapture is not supported in ROCm')
     # TODO(leofang): check and raise if stream == 0?
     with nogil:
         status = cudaStreamBeginCapture(<driver.Stream>stream,
@@ -945,8 +937,6 @@ cpdef streamBeginCapture(intptr_t stream, int mode=streamCaptureModeRelaxed):
 cpdef intptr_t streamEndCapture(intptr_t stream) except? 0:
     # TODO(leofang): check and raise if stream == 0?
     cdef Graph g
-    if _is_hip_environment:
-        raise RuntimeError('streamEndCapture is not supported in ROCm')
     with nogil:
         status = cudaStreamEndCapture(<driver.Stream>stream, &g)
     check_status(status)
@@ -955,8 +945,6 @@ cpdef intptr_t streamEndCapture(intptr_t stream) except? 0:
 
 cpdef bint streamIsCapturing(intptr_t stream) except*:
     cdef StreamCaptureStatus s
-    if _is_hip_environment:
-        raise RuntimeError('streamIsCapturing is not supported in ROCm')
     with nogil:
         status = cudaStreamIsCapturing(<driver.Stream>stream, &s)
     check_status(status)  # cudaErrorStreamCaptureImplicit could be raised here

cupy_backends/cuda/libs/cublas.pyx

@@ -508,7 +508,7 @@ cpdef setStream(intptr_t handle, size_t stream):
     # https://docs.nvidia.com/cuda/cublas/index.html#CUDA-graphs
     # Before we come up with a robust strategy to test the support conditions,
     # we disable this functionality.
-    if not runtime._is_hip_environment and runtime.streamIsCapturing(stream):
+    if runtime.streamIsCapturing(stream):
         raise NotImplementedError(
             'calling cuBLAS API during stream capture is currently '
             'unsupported')

cupy_backends/cuda/libs/curand.pyx

@@ -115,7 +115,7 @@ cpdef int getVersion() except? -1:
 cpdef setStream(size_t generator, size_t stream):
     # TODO(leofang): The support of stream capture is not mentioned at all in
     # the cuRAND docs (as of CUDA 11.5), so we disable this functionality.
-    if not runtime._is_hip_environment and runtime.streamIsCapturing(stream):
+    if runtime.streamIsCapturing(stream):
         raise NotImplementedError(
             'calling cuRAND API during stream capture is currently '
             'unsupported')

cupy_backends/cuda/libs/cusolver.pyx

@@ -1136,7 +1136,7 @@ cpdef spDestroy(intptr_t handle):
 cpdef setStream(intptr_t handle, size_t stream):
     # TODO(leofang): The support of stream capture is not mentioned at all in
     # the cuSOLVER docs (as of CUDA 11.5), so we disable this functionality.
-    if not runtime._is_hip_environment and runtime.streamIsCapturing(stream):
+    if runtime.streamIsCapturing(stream):
         raise NotImplementedError(
             'calling cuSOLVER API during stream capture is currently '
             'unsupported')

cupy_backends/cuda/libs/cusparse.pyx

@@ -1630,7 +1630,7 @@ cpdef void setStream(intptr_t handle, size_t stream) except *:
     # https://docs.nvidia.com/cuda/cusparse/index.html#optimization-notes
     # Before we come up with a robust strategy to test the support conditions,
     # we disable this functionality.
-    if not runtime._is_hip_environment and runtime.streamIsCapturing(stream):
+    if runtime.streamIsCapturing(stream):
         raise NotImplementedError(
             'calling cuSPARSE API during stream capture is currently '
             'unsupported')

cupy_backends/hip/cupy_hip_common.h

@@ -2,8 +2,13 @@
 #define INCLUDE_GUARD_HIP_CUPY_COMMON_H
 
 #include <hip/hip_runtime_api.h>
+#if HIP_VERSION >= 50530600
+#include <hipblas/hipblas.h>
+#include <rocsolver/rocsolver.h>
+#else
 #include <hipblas.h>
 #include <rocsolver.h>
+#endif
 
 #define CUDA_VERSION 0