Replace use of custom CUDA bindings with CUDA-Python

python/rmm/_cuda/gpu.pyx → python/rmm/_cuda/gpu.py

@@ -1,59 +1,36 @@
 # Copyright (c) 2020, NVIDIA CORPORATION.
 
-from rmm._cuda.gpu cimport (
-    CUresult,
-    cudaDeviceAttr,
-    cudaDeviceGetAttribute,
-    cudaDeviceProp,
-    cudaDriverGetVersion,
-    cudaError,
-    cudaError_t,
-    cudaGetDeviceCount,
-    cudaGetDeviceProperties,
-    cudaGetErrorName,
-    cudaGetErrorString,
-    cudaRuntimeGetVersion,
-    cuDeviceGetName,
-    cuGetErrorName,
-    cuGetErrorString,
-)
-
-from enum import IntEnum
+from cuda import cuda, cudart
 
 
 class CUDARuntimeError(RuntimeError):
-
-    def __init__(self, cudaError_t status):
+    def __init__(self, status: cuda.CUresult.CUDA_SUCCESS):
         self.status = status
-        cdef str name = cudaGetErrorName(status).decode()
-        cdef str msg = cudaGetErrorString(status).decode()
+        _, name = cudart.cudaGetErrorName(status)
+        _, msg = cudart.cudaGetErrorString(status)
         super(CUDARuntimeError, self).__init__(
-            '%s: %s' % (name, msg))
+            "%s: %s" % (name.decode(), msg.decode())
+        )
 
     def __reduce__(self):
         return (type(self), (self.status,))
 
 
 class CUDADriverError(RuntimeError):
-
-    def __init__(self, CUresult status):
+    def __init__(self, status: cuda.CUresult):
         self.status = status
 
-        cdef const char* name_cstr
-        cdef CUresult name_status = cuGetErrorName(status, &name_cstr)
-        if name_status != 0:
-            raise CUDADriverError(name_status)
-
-        cdef const char* msg_cstr
-        cdef CUresult msg_status = cuGetErrorString(status, &msg_cstr)
-        if msg_status != 0:
-            raise CUDADriverError(msg_status)
+        err, name = cuda.cuGetErrorName(status)
+        if err != cuda.CUresult.CUDA_SUCCESS:
+            raise CUDADriverError(err.value)
 
-        cdef str name = name_cstr.decode()
-        cdef str msg = msg_cstr.decode()
+        err, msg = cuda.cuGetErrorString(status)
+        if err != cuda.CUresult.CUDA_SUCCESS:
+            raise CUDADriverError(err.value)
 
         super(CUDADriverError, self).__init__(
-            '%s: %s' % (name, msg))
+            "%s: %s" % (name.decode(), msg.decode())
+        )
 
     def __reduce__(self):
         return (type(self), (self.status,))
@@ -69,9 +46,8 @@ def driverGetVersion():
     This function automatically raises CUDARuntimeError with error message
     and status code.
     """
-    cdef int version
-    cdef cudaError_t status = cudaDriverGetVersion(&version)
-    if status != cudaError.cudaSuccess:
+    status, version = cudart.cudaDriverGetVersion()
+    if status != cudart.cudaError_t.cudaSuccess:
         raise CUDARuntimeError(status)
     return version
 
@@ -80,24 +56,22 @@ def getDevice():
     """
     Get the current CUDA device
     """
-    cdef int current_device
-    cdef cudaError_t status = cudaGetDevice(&current_device)
-    if status != cudaError.cudaSuccess:
+    status, device = cudart.cudaGetDevice()
+    if status != cudart.cudaError_t.cudaSuccess:
         raise CUDARuntimeError(status)
-    return current_device
+    return device
 
 
-def setDevice(int device):
+def setDevice(device: int):
     """
     Set the current CUDA device
     Parameters
     ----------
     device : int
         The ID of the device to set as current
     """
-    cdef cudaError_t status = cudaSetDevice(device)
-
-    if status != cudaError.cudaSuccess:
+    (status,) = cudart.cudaSetDevice(device)
+    if status != cudart.cudaError_t.cudaSuccess:
         raise CUDARuntimeError(status)
 
 
@@ -110,10 +84,8 @@ def runtimeGetVersion():
     This function automatically raises CUDARuntimeError with error message
     and status code.
     """
-
-    cdef int version
-    cdef cudaError_t status = cudaRuntimeGetVersion(&version)
-    if status != cudaError.cudaSuccess:
+    status, version = cudart.cudaRuntimeGetVersion()
+    if status != cudart.cudaError_t.cudaSuccess:
         raise CUDARuntimeError(status)
     return version
 
@@ -126,16 +98,13 @@ def getDeviceCount():
     This function automatically raises CUDARuntimeError with error message
     and status code.
     """
-
-    cdef int count
-    cdef cudaError_t status = cudaGetDeviceCount(&count)
-
-    if status != cudaError.cudaSuccess:
+    status, count = cudart.cudaGetDeviceCount()
+    if status != cudart.cudaError_t.cudaSuccess:
         raise CUDARuntimeError(status)
     return count
 
 
-def getDeviceAttribute(cudaDeviceAttr attr, int device):
+def getDeviceAttribute(attr: cudart.cudaDeviceAttr, device: int):
     """
     Returns information about the device.
 
@@ -149,15 +118,13 @@ def getDeviceAttribute(cudaDeviceAttr attr, int device):
     This function automatically raises CUDARuntimeError with error message
     and status code.
     """
-
-    cdef int value
-    cdef cudaError_t status = cudaDeviceGetAttribute(&value, attr, device)
-    if status != cudaError.cudaSuccess:
+    status, value = cudart.cudaDeviceGetAttribute(attr, device)
+    if status != cudart.cudaError_t.cudaSuccess:
         raise CUDARuntimeError(status)
     return value
 
 
-def getDeviceProperties(int device):
+def getDeviceProperties(device: int):
     """
     Returns information about the compute-device.
 
@@ -169,15 +136,13 @@ def getDeviceProperties(int device):
     This function automatically raises CUDARuntimeError with error message
     and status code.
     """
-
-    cdef cudaDeviceProp prop
-    cdef cudaError_t status = cudaGetDeviceProperties(&prop, device)
-    if status != cudaError.cudaSuccess:
+    status, prop = cudart.cudaGetDeviceProperties(device)
+    if status != cudart.cudaError_t.cudaSuccess:
         raise CUDARuntimeError(status)
     return prop
 
 
-def deviceGetName(int device):
+def deviceGetName(device: int):
     """
     Returns an identifer string for the device.
 
@@ -190,12 +155,7 @@ def deviceGetName(int device):
     and status code.
     """
 
-    cdef char[256] device_name
-    cdef CUresult status = cuDeviceGetName(
-        device_name,
-        sizeof(device_name),
-        device
-    )
-    if status != 0:
+    status, device_name = cuda.cuDeviceGetName(256, cuda.CUdevice(device))
+    if status.value != 0:
         raise CUDADriverError(status)
     return device_name.decode()

python/rmm/_cuda/stream.pxd

@@ -12,11 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from cuda.ccudart cimport cudaStream_t
 from libc.stdint cimport uintptr_t
 from libcpp cimport bool
 
 from rmm._lib.cuda_stream_view cimport cuda_stream_view
-from rmm._lib.lib cimport cudaStream_t
 
 
 cdef class Stream:

python/rmm/_cuda/stream.pyx

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from cuda.ccudart cimport cudaStream_t
 from libc.stdint cimport uintptr_t
 from libcpp cimport bool
 
@@ -21,7 +22,6 @@ from rmm._lib.cuda_stream_view cimport (
     cuda_stream_per_thread,
     cuda_stream_view,
 )
-from rmm._lib.lib cimport cudaStream_t
 
 from numba import cuda
 

python/rmm/_lib/cuda_stream.pxd

@@ -13,11 +13,11 @@
 # limitations under the License.
 
 cimport cython
+from cuda.ccudart cimport cudaStream_t
 from libcpp cimport bool
 from libcpp.memory cimport unique_ptr
 
 from rmm._lib.cuda_stream_view cimport cuda_stream_view
-from rmm._lib.lib cimport cudaStream_t
 
 
 cdef extern from "rmm/cuda_stream.hpp" namespace "rmm" nogil:

python/rmm/_lib/cuda_stream.pyx

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 cimport cython
+from cuda.ccudart cimport cudaStream_t
 from libc.stdint cimport uintptr_t
 from libcpp cimport bool
 

python/rmm/_lib/cuda_stream_view.pxd

@@ -12,10 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from cuda.ccudart cimport cudaStream_t
 from libcpp cimport bool
 
-from rmm._lib.lib cimport cudaStream_t
-
 
 cdef extern from "rmm/cuda_stream_view.hpp" namespace "rmm" nogil:
     cdef cppclass cuda_stream_view:

python/rmm/_lib/device_buffer.pyx

@@ -20,20 +20,20 @@ from libc.stdint cimport uintptr_t
 from libcpp.memory cimport unique_ptr
 from libcpp.utility cimport move
 
-from rmm._cuda.gpu cimport cudaError, cudaError_t
 from rmm._cuda.stream cimport Stream
 
 from rmm._cuda.stream import DEFAULT_STREAM
 
-from rmm._lib.lib cimport (
+cimport cuda.ccudart as ccudart
+from cuda.ccudart cimport (
+    cudaError,
+    cudaError_t,
     cudaMemcpyAsync,
-    cudaMemcpyDeviceToDevice,
-    cudaMemcpyDeviceToHost,
-    cudaMemcpyHostToDevice,
     cudaMemcpyKind,
     cudaStream_t,
     cudaStreamSynchronize,
 )
+
 from rmm._lib.memory_resource cimport get_current_device_resource
 
 
@@ -73,7 +73,6 @@ cdef class DeviceBuffer:
         >>> db = rmm.DeviceBuffer(size=5)
         """
         cdef const void* c_ptr
-        cdef cudaError_t err
 
         with nogil:
             c_ptr = <const void*>ptr
@@ -339,7 +338,7 @@ cpdef DeviceBuffer to_device(const unsigned char[::1] b,
 cdef void _copy_async(const void* src,
                       void* dst,
                       size_t count,
-                      cudaMemcpyKind kind,
+                      ccudart.cudaMemcpyKind kind,
                       cuda_stream_view stream) nogil:
     """
     Asynchronously copy data between host and/or device pointers
@@ -398,7 +397,7 @@ cpdef void copy_ptr_to_host(uintptr_t db,
 
     with nogil:
         _copy_async(<const void*>db, <void*>&hb[0], len(hb),
-                    cudaMemcpyDeviceToHost, stream.view())
+                    cudaMemcpyKind.cudaMemcpyDeviceToHost, stream.view())
 
     if stream.c_is_default():
         stream.c_synchronize()
@@ -442,7 +441,7 @@ cpdef void copy_host_to_ptr(const unsigned char[::1] hb,
 
     with nogil:
         _copy_async(<const void*>&hb[0], <void*>db, len(hb),
-                    cudaMemcpyHostToDevice, stream.view())
+                    cudaMemcpyKind.cudaMemcpyHostToDevice, stream.view())
 
     if stream.c_is_default():
         stream.c_synchronize()
@@ -475,4 +474,4 @@ cpdef void copy_device_to_ptr(uintptr_t d_src,
 
     with nogil:
         _copy_async(<const void*>d_src, <void*>d_dst, count,
-                    cudaMemcpyDeviceToDevice, stream.view())
+                    cudaMemcpyKind.cudaMemcpyDeviceToDevice, stream.view())

python/rmm/_lib/lib.pxd

@@ -17,24 +17,4 @@ from libcpp cimport bool
 from libcpp.utility cimport pair
 from libcpp.vector cimport vector
 
-from rmm._cuda.gpu cimport cudaError_t
-
 ctypedef pair[const char*, unsigned int] caller_pair
-
-
-cdef extern from * nogil:
-
-    ctypedef void* cudaStream_t "cudaStream_t"
-
-    ctypedef enum cudaMemcpyKind "cudaMemcpyKind":
-        cudaMemcpyHostToHost = 0
-        cudaMemcpyHostToDevice = 1
-        cudaMemcpyDeviceToHost = 2
-        cudaMemcpyDeviceToDevice = 3
-        cudaMemcpyDefault = 4
-
-    cudaError_t cudaMemcpyAsync(void* dst, const void* src, size_t count,
-                                cudaMemcpyKind kind)
-    cudaError_t cudaMemcpyAsync(void* dst, const void* src, size_t count,
-                                cudaMemcpyKind kind, cudaStream_t stream)
-    cudaError_t cudaStreamSynchronize(cudaStream_t stream)

python/rmm/_lib/memory_resource.pxd

@@ -68,7 +68,4 @@ cdef class StatisticsResourceAdaptor(UpstreamResourceAdaptor):
 cdef class TrackingResourceAdaptor(UpstreamResourceAdaptor):
     pass
 
-cdef class FailureCallbackResourceAdaptor(UpstreamResourceAdaptor):
-    cdef object _callback
-
 cpdef DeviceMemoryResource get_current_device_resource()

python/rmm/_lib/memory_resource.pyx

@@ -23,7 +23,9 @@ from libcpp.cast cimport dynamic_cast
 from libcpp.memory cimport make_shared, make_unique, shared_ptr, unique_ptr
 from libcpp.string cimport string
 
-from rmm._cuda.gpu import CUDARuntimeError, cudaError, getDevice, setDevice
+from cuda.cudart import cudaError_t
+
+from rmm._cuda.gpu import CUDARuntimeError, getDevice, setDevice
 
 
 # NOTE: Keep extern declarations in .pyx file as much as possible to avoid
@@ -705,7 +707,7 @@ cpdef void _initialize(
     try:
         original_device = getDevice()
     except CUDARuntimeError as e:
-        if e.status == cudaError.cudaErrorNoDevice:
+        if e.status == cudaError_t.cudaErrorNoDevice:
             warnings.warn(e.msg)
         else:
             raise e