Enable NV bindings by default and add them as a dependency (#284)

* use nv binding by default and add dep

* docs

* address review

* move code around

* reset file

* generic cuda-core conda dep

* fix

* fix logic

* update

* update deps

* updates

* Update numba_cuda/numba/cuda/cudadrv/driver.py

Co-authored-by: Graham Markall <[email protected]>

* use minimal bindings dep

---------

Co-authored-by: Graham Markall <[email protected]>

Author: brandon-b-miller

conda/recipes/numba-cuda/meta.yaml

@@ -28,6 +28,7 @@ requirements:
   run:
     - python
     - numba >=0.59.1
+    - cuda-bindings
 
 about:
   home: {{ project_urls["Homepage"] }}

docs/source/reference/envvars.rst

@@ -108,9 +108,8 @@ target.
 
    When set to 1, Numba will attempt to use the `NVIDIA CUDA Python binding
    <https://nvidia.github.io/cuda-python/>`_ to make calls to the driver API
-   instead of using its own ctypes binding. This defaults to 0 (off), as the
-   NVIDIA binding is currently missing support for Per-Thread Default
-   Streams and the profiler APIs.
+   instead of using its own ctypes binding. This defaults to 1 (on). Set to
+   0 to use the ctypes bindings.
 
 .. envvar:: NUMBA_CUDA_INCLUDE_PATH
 

numba_cuda/numba/cuda/__init__.py

@@ -1,5 +1,36 @@
+import importlib
 from numba import runtests
 from numba.core import config
+from .utils import _readenv
+
+# Enable pynvjitlink if the environment variables NUMBA_CUDA_ENABLE_PYNVJITLINK
+# or CUDA_ENABLE_PYNVJITLINK are set, or if the pynvjitlink module is found. If
+# explicitly disabled, do not use pynvjitlink, even if present in the env.
+_pynvjitlink_enabled_in_env = _readenv(
+    "NUMBA_CUDA_ENABLE_PYNVJITLINK", bool, None
+)
+_pynvjitlink_enabled_in_cfg = getattr(config, "CUDA_ENABLE_PYNVJITLINK", None)
+
+if _pynvjitlink_enabled_in_env is not None:
+    ENABLE_PYNVJITLINK = _pynvjitlink_enabled_in_env
+elif _pynvjitlink_enabled_in_cfg is not None:
+    ENABLE_PYNVJITLINK = _pynvjitlink_enabled_in_cfg
+else:
+    ENABLE_PYNVJITLINK = importlib.util.find_spec("pynvjitlink") is not None
+
+if not hasattr(config, "CUDA_ENABLE_PYNVJITLINK"):
+    config.CUDA_ENABLE_PYNVJITLINK = ENABLE_PYNVJITLINK
+
+# Upstream numba sets CUDA_USE_NVIDIA_BINDING to 0 by default, so it always
+# exists. Override, but not if explicitly set to 0 in the envioronment.
+_nvidia_binding_enabled_in_env = _readenv(
+    "NUMBA_CUDA_USE_NVIDIA_BINDING", bool, None
+)
+if _nvidia_binding_enabled_in_env is False:
+    USE_NV_BINDING = False
+else:
+    USE_NV_BINDING = True
+    config.CUDA_USE_NVIDIA_BINDING = USE_NV_BINDING
 
 if config.ENABLE_CUDASIM:
     from .simulator_init import *

numba_cuda/numba/cuda/cudadrv/driver.py

@@ -49,23 +49,14 @@
 from .drvapi import cu_occupancy_b2d_size, cu_stream_callback_pyobj, cu_uuid
 from .mappings import FILE_EXTENSION_MAP
 from .linkable_code import LinkableCode, LTOIR, Fatbin, Object
-from numba.cuda.utils import _readenv, cached_file_read
+from numba.cuda.utils import cached_file_read
 from numba.cuda.cudadrv import enums, drvapi, nvrtc
 
 try:
     from pynvjitlink.api import NvJitLinker, NvJitLinkError
 except ImportError:
     NvJitLinker, NvJitLinkError = None, None
 
-USE_NV_BINDING = config.CUDA_USE_NVIDIA_BINDING
-
-if USE_NV_BINDING:
-    from cuda import cuda as binding
-
-    # There is no definition of the default stream in the Nvidia bindings (nor
-    # is there at the C/C++ level), so we define it here so we don't need to
-    # use a magic number 0 in places where we want the default stream.
-    CU_STREAM_DEFAULT = 0
 
 MIN_REQUIRED_CC = (3, 5)
 SUPPORTS_IPC = sys.platform.startswith("linux")
@@ -82,23 +73,15 @@
     "to be available"
 )
 
-# Enable pynvjitlink if the environment variables NUMBA_CUDA_ENABLE_PYNVJITLINK
-# or CUDA_ENABLE_PYNVJITLINK are set, or if the pynvjitlink module is found. If
-# explicitly disabled, do not use pynvjitlink, even if present in the env.
-_pynvjitlink_enabled_in_env = _readenv(
-    "NUMBA_CUDA_ENABLE_PYNVJITLINK", bool, None
-)
-_pynvjitlink_enabled_in_cfg = getattr(config, "CUDA_ENABLE_PYNVJITLINK", None)
+USE_NV_BINDING = config.CUDA_USE_NVIDIA_BINDING
 
-if _pynvjitlink_enabled_in_env is not None:
-    ENABLE_PYNVJITLINK = _pynvjitlink_enabled_in_env
-elif _pynvjitlink_enabled_in_cfg is not None:
-    ENABLE_PYNVJITLINK = _pynvjitlink_enabled_in_cfg
-else:
-    ENABLE_PYNVJITLINK = importlib.util.find_spec("pynvjitlink") is not None
+if USE_NV_BINDING:
+    from cuda.bindings import driver as binding
 
-if not hasattr(config, "CUDA_ENABLE_PYNVJITLINK"):
-    config.CUDA_ENABLE_PYNVJITLINK = ENABLE_PYNVJITLINK
+    # There is no definition of the default stream in the Nvidia bindings (nor
+    # is there at the C/C++ level), so we define it here so we don't need to
+    # use a magic number 0 in places where we want the default stream.
+    CU_STREAM_DEFAULT = 0
 
 
 def make_logger():
@@ -3192,7 +3175,6 @@ def __init__(self, max_registers=0, lineinfo=False, cc=None):
 
         raw_keys = list(options.keys())
         raw_values = list(options.values())
-
         self.handle = driver.cuLinkCreate(len(raw_keys), raw_keys, raw_values)
 
         weakref.finalize(self, driver.cuLinkDestroy, self.handle)

numba_cuda/numba/cuda/cudadrv/mappings.py

@@ -2,9 +2,9 @@
 from . import enums
 
 if config.CUDA_USE_NVIDIA_BINDING:
-    from cuda import cuda
+    from cuda.bindings import driver
 
-    jitty = cuda.CUjitInputType
+    jitty = driver.CUjitInputType
     FILE_EXTENSION_MAP = {
         "o": jitty.CU_JIT_INPUT_OBJECT,
         "ptx": jitty.CU_JIT_INPUT_PTX,

numba_cuda/numba/cuda/tests/cudadrv/test_linker.py

@@ -104,8 +104,6 @@ def simple_lmem(A, B, dty):
 
 @skip_on_cudasim("Linking unsupported in the simulator")
 class TestLinker(CUDATestCase):
-    _NUMBA_NVIDIA_BINDING_0_ENV = {"NUMBA_CUDA_USE_NVIDIA_BINDING": "0"}
-
     @require_context
     def test_linker_basic(self):
         """Simply go through the constructor and destructor"""

numba_cuda/numba/cuda/tests/cudadrv/test_module_callbacks.py

@@ -15,7 +15,7 @@
     from cuda.bindings.driver import cuModuleGetGlobal, cuMemcpyHtoD
 
     if config.CUDA_USE_NVIDIA_BINDING:
-        from cuda.cuda import CUmodule as cu_module_type
+        from cuda.bindings.driver import CUmodule as cu_module_type
     else:
         from numba.cuda.cudadrv.drvapi import cu_module as cu_module_type
 

numba_cuda/numba/cuda/tests/cudadrv/test_nvjitlink.py

@@ -57,8 +57,6 @@
 )
 @skip_on_cudasim("Linking unsupported in the simulator")
 class TestLinker(CUDATestCase):
-    _NUMBA_NVIDIA_BINDING_0_ENV = {"NUMBA_CUDA_USE_NVIDIA_BINDING": "0"}
-
     def test_nvjitlink_create(self):
         patched_linker = PyNvJitLinker(cc=(7, 5))
         assert "-arch=sm_75" in patched_linker.options

pyproject.toml

@@ -22,12 +22,14 @@ dependencies = ["numba>=0.59.1"]
 
 [project.optional-dependencies]
 cu11 = [
+    "cuda-bindings==11.8.*",
     "cuda-python==11.8.*",  # supports all CTK 11.x
     "nvidia-cuda-nvcc-cu11",  # for libNVVM
     "nvidia-cuda-runtime-cu11",
     "nvidia-cuda-nvrtc-cu11",
 ]
 cu12 = [
+    "cuda-bindings==12.9.*",
     "cuda-python==12.9.*",  # supports all CTK 12.x
     "nvidia-cuda-nvcc-cu12",  # for libNVVM
     "nvidia-cuda-runtime-cu12",
@@ -110,7 +112,7 @@ exclude = [
 # Slightly long line in the standard version file
 "numba_cuda/_version.py" = ["E501"]
 # "Unused" imports / potentially undefined names in init files
-"numba_cuda/numba/cuda/__init__.py" = ["F401", "F403", "F405"]
+"numba_cuda/numba/cuda/__init__.py" = ["F401", "F403", "F405", "E402"]
 "numba_cuda/numba/cuda/simulator/__init__.py" = ["F401", "F403"]
 "numba_cuda/numba/cuda/simulator/cudadrv/__init__.py" = ["F401"]
 # Ignore star imports", " unused imports", " and "may be defined by star imports"