Use bindings layout for all cuda-python imports. (#1756)

Closes #1755.

Replaces #1675.

Authors:
  - Bradley Dice (https://github.com/bdice)
  - Matthew Murray (https://github.com/Matt711)

Approvers:
  - Matthew Murray (https://github.com/Matt711)
  - Vyas Ramasubramani (https://github.com/vyasr)

URL: #1756

python/rmm/rmm/_cuda/gpu.py

@@ -1,14 +1,14 @@
-# Copyright (c) 2020, NVIDIA CORPORATION.
+# Copyright (c) 2020-2024, NVIDIA CORPORATION.
 
-from cuda import cuda, cudart
+from cuda.bindings import driver, runtime
 
 
 class CUDARuntimeError(RuntimeError):
-    def __init__(self, status: cudart.cudaError_t):
+    def __init__(self, status: runtime.cudaError_t):
         self.status = status
 
-        _, name = cudart.cudaGetErrorName(status)
-        _, msg = cudart.cudaGetErrorString(status)
+        _, name = runtime.cudaGetErrorName(status)
+        _, msg = runtime.cudaGetErrorString(status)
 
         super(CUDARuntimeError, self).__init__(
             f"{name.decode()}: {msg.decode()}"
@@ -19,11 +19,11 @@ def __reduce__(self):
 
 
 class CUDADriverError(RuntimeError):
-    def __init__(self, status: cuda.CUresult):
+    def __init__(self, status: driver.CUresult):
         self.status = status
 
-        _, name = cuda.cuGetErrorName(status)
-        _, msg = cuda.cuGetErrorString(status)
+        _, name = driver.cuGetErrorName(status)
+        _, msg = driver.cuGetErrorString(status)
 
         super(CUDADriverError, self).__init__(
             f"{name.decode()}: {msg.decode()}"
@@ -43,8 +43,8 @@ def driverGetVersion():
     This function automatically raises CUDARuntimeError with error message
     and status code.
     """
-    status, version = cudart.cudaDriverGetVersion()
-    if status != cudart.cudaError_t.cudaSuccess:
+    status, version = runtime.cudaDriverGetVersion()
+    if status != runtime.cudaError_t.cudaSuccess:
         raise CUDARuntimeError(status)
     return version
 
@@ -53,8 +53,8 @@ def getDevice():
     """
     Get the current CUDA device
     """
-    status, device = cudart.cudaGetDevice()
-    if status != cudart.cudaError_t.cudaSuccess:
+    status, device = runtime.cudaGetDevice()
+    if status != runtime.cudaError_t.cudaSuccess:
         raise CUDARuntimeError(status)
     return device
 
@@ -67,26 +67,25 @@ def setDevice(device: int):
     device : int
         The ID of the device to set as current
     """
-    (status,) = cudart.cudaSetDevice(device)
-    if status != cudart.cudaError_t.cudaSuccess:
+    (status,) = runtime.cudaSetDevice(device)
+    if status != runtime.cudaError_t.cudaSuccess:
         raise CUDARuntimeError(status)
 
 
 def runtimeGetVersion():
     """
-    Returns the version number of the current CUDA Runtime instance.
-    The version is returned as (1000 major + 10 minor). For example,
-    CUDA 9.2 would be represented by 9020.
+    Returns the version number of the local CUDA runtime.
 
-    This calls numba.cuda.runtime.get_version() rather than cuda-python due to
-    current limitations in cuda-python.
-    """
-    # TODO: Replace this with `cuda.cudart.cudaRuntimeGetVersion()` when the
-    # limitation is fixed.
-    import numba.cuda
+    The version is returned as ``(1000 * major + 10 * minor)``. For example,
+    CUDA 12.5 would be represented by 12050.
 
-    major, minor = numba.cuda.runtime.get_version()
-    return major * 1000 + minor * 10
+    This function automatically raises CUDARuntimeError with error message
+    and status code.
+    """
+    status, version = runtime.getLocalRuntimeVersion()
+    if status != runtime.cudaError_t.cudaSuccess:
+        raise CUDARuntimeError(status)
+    return version
 
 
 def getDeviceCount():
@@ -97,13 +96,13 @@ def getDeviceCount():
     This function automatically raises CUDARuntimeError with error message
     and status code.
     """
-    status, count = cudart.cudaGetDeviceCount()
-    if status != cudart.cudaError_t.cudaSuccess:
+    status, count = runtime.cudaGetDeviceCount()
+    if status != runtime.cudaError_t.cudaSuccess:
         raise CUDARuntimeError(status)
     return count
 
 
-def getDeviceAttribute(attr: cudart.cudaDeviceAttr, device: int):
+def getDeviceAttribute(attr: runtime.cudaDeviceAttr, device: int):
     """
     Returns information about the device.
 
@@ -117,8 +116,8 @@ def getDeviceAttribute(attr: cudart.cudaDeviceAttr, device: int):
     This function automatically raises CUDARuntimeError with error message
     and status code.
     """
-    status, value = cudart.cudaDeviceGetAttribute(attr, device)
-    if status != cudart.cudaError_t.cudaSuccess:
+    status, value = runtime.cudaDeviceGetAttribute(attr, device)
+    if status != runtime.cudaError_t.cudaSuccess:
         raise CUDARuntimeError(status)
     return value
 
@@ -135,8 +134,8 @@ def getDeviceProperties(device: int):
     This function automatically raises CUDARuntimeError with error message
     and status code.
     """
-    status, prop = cudart.cudaGetDeviceProperties(device)
-    if status != cudart.cudaError_t.cudaSuccess:
+    status, prop = runtime.cudaGetDeviceProperties(device)
+    if status != runtime.cudaError_t.cudaSuccess:
         raise CUDARuntimeError(status)
     return prop
 
@@ -154,7 +153,7 @@ def deviceGetName(device: int):
     and status code.
     """
 
-    status, device_name = cuda.cuDeviceGetName(256, cuda.CUdevice(device))
-    if status != cuda.CUresult.CUDA_SUCCESS:
+    status, device_name = driver.cuDeviceGetName(256, driver.CUdevice(device))
+    if status != driver.CUresult.CUDA_SUCCESS:
         raise CUDADriverError(status)
     return device_name.decode()

python/rmm/rmm/_cuda/stream.pxd

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from cuda.ccudart cimport cudaStream_t
+from cuda.bindings.cyruntime cimport cudaStream_t
 from libc.stdint cimport uintptr_t
 from libcpp cimport bool
 

python/rmm/rmm/_cuda/stream.pyx

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from cuda.ccudart cimport cudaStream_t
+from cuda.bindings.cyruntime cimport cudaStream_t
 from libc.stdint cimport uintptr_t
 from libcpp cimport bool
 

python/rmm/rmm/allocators/numba.py

@@ -15,7 +15,7 @@
 import ctypes
 import inspect
 
-from cuda.cuda import CUdeviceptr, cuIpcGetMemHandle
+from cuda.bindings.driver import CUdeviceptr, cuIpcGetMemHandle
 from numba import config, cuda
 from numba.cuda import HostOnlyCUDAMemoryManager, IpcHandle, MemoryPointer
 

python/rmm/rmm/librmm/cuda_stream.pxd

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from cuda.ccudart cimport cudaStream_t
+from cuda.bindings.cyruntime cimport cudaStream_t
 from libcpp cimport bool
 
 from rmm.librmm.cuda_stream_view cimport cuda_stream_view

python/rmm/rmm/librmm/cuda_stream_view.pxd

@@ -1,4 +1,4 @@
-# Copyright (c) 2020, NVIDIA CORPORATION.
+# Copyright (c) 2020-2024, NVIDIA CORPORATION.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from cuda.ccudart cimport cudaStream_t
+from cuda.bindings.cyruntime cimport cudaStream_t
 from libcpp cimport bool
 
 

python/rmm/rmm/pylibrmm/cuda_stream.pxd

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 cimport cython
-from cuda.ccudart cimport cudaStream_t
+from cuda.bindings.cyruntime cimport cudaStream_t
 from libcpp cimport bool
 from libcpp.memory cimport unique_ptr
 

python/rmm/rmm/pylibrmm/cuda_stream.pyx

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 cimport cython
-from cuda.ccudart cimport cudaStream_t
+from cuda.bindings.cyruntime cimport cudaStream_t
 from libcpp cimport bool
 
 from rmm.librmm.cuda_stream cimport cuda_stream

python/rmm/rmm/pylibrmm/device_buffer.pyx

@@ -23,8 +23,7 @@ from rmm._cuda.stream cimport Stream
 
 from rmm._cuda.stream import DEFAULT_STREAM
 
-cimport cuda.ccudart as ccudart
-from cuda.ccudart cimport (
+from cuda.bindings.cyruntime cimport (
     cudaError,
     cudaError_t,
     cudaMemcpyAsync,
@@ -421,7 +420,7 @@ cpdef DeviceBuffer to_device(const unsigned char[::1] b,
 cdef void _copy_async(const void* src,
                       void* dst,
                       size_t count,
-                      ccudart.cudaMemcpyKind kind,
+                      cudaMemcpyKind kind,
                       cuda_stream_view stream) except * nogil:
     """
     Asynchronously copy data between host and/or device pointers.

python/rmm/rmm/pylibrmm/memory_resource.pyx

@@ -28,7 +28,7 @@ from libcpp.memory cimport make_unique, unique_ptr
 from libcpp.optional cimport optional
 from libcpp.pair cimport pair
 
-from cuda.cudart import cudaError_t
+from cuda.bindings.runtime import cudaError_t
 
 from rmm._cuda.gpu import CUDARuntimeError, getDevice, setDevice