fallback to cuFuncAPI

Author: ksimpson-work

cuda_core/cuda/core/experimental/_module.py

@@ -3,13 +3,16 @@
 # SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE
 
 
+from warnings import warn
+
 from cuda.core.experimental._utils import driver, get_binding_version, handle_return, precondition
 
 _backend = {
     "old": {
         "file": driver.cuModuleLoad,
         "data": driver.cuModuleLoadDataEx,
         "kernel": driver.cuModuleGetFunction,
+        "attribute": driver.cuFuncGetAttribute,
     },
 }
 
@@ -34,6 +37,7 @@ def _lazy_init():
             "file": driver.cuLibraryLoadFromFile,
             "data": driver.cuLibraryLoadData,
             "kernel": driver.cuLibraryGetKernel,
+            "attribute": driver.cuKernelGetAttribute,
         }
         _kernel_ctypes = (driver.CUfunction, driver.CUkernel)
     else:
@@ -46,19 +50,30 @@ class KernelAttributes:
     def __init__(self):
         raise RuntimeError("KernelAttributes should not be instantiated directly")
 
-    slots = ("_handle", "_cache")
+    slots = ("_handle", "_cache", "_backend_version", "_loader")
 
     def _init(handle):
         self = KernelAttributes.__new__(KernelAttributes)
         self._handle = handle
         self._cache = {}
+
+        self._backend_version = "new" if (_py_major_ver >= 12 and _driver_ver >= 12000) else "old"
+        self._loader = _backend[self._backend_version]
         return self
 
     def _get_cached_attribute(self, device_id: int, attribute: driver.CUfunction_attribute) -> int:
         """Helper function to get a cached attribute or fetch and cache it if not present."""
         if device_id in self._cache and attribute in self._cache[device_id]:
             return self._cache[device_id][attribute]
-        result = handle_return(driver.cuKernelGetAttribute(attribute, self._handle, device_id))
+        if self._backend_version == "new":
+            result = handle_return(self._loader["attribute"](attribute, self._handle, device_id))
+        else:  # "old" backend
+            warn(
+                "Device ID argument is ignored when getting attribute from kernel when cuda version < 12. ",
+                RuntimeWarning,
+                stacklevel=2,
+            )
+            result = handle_return(self._loader["attribute"](attribute, self._handle))
         if device_id not in self._cache:
             self._cache[device_id] = {}
         self._cache[device_id][attribute] = result

cuda_core/tests/test_module.py

@@ -10,21 +10,7 @@
 import pytest
 from conftest import can_load_generated_ptx
 
-try:
-    from cuda.bindings import driver
-except ImportError:
-    from cuda import cuda as driver
-
 from cuda.core.experimental import Program, ProgramOptions, system
-from cuda.core.experimental._utils import get_binding_version, handle_return
-
-
-@pytest.fixture(scope="module")
-def cuda_version():
-    # binding availability depends on cuda-python version
-    _py_major_ver, _ = get_binding_version()
-    _driver_ver = handle_return(driver.cuDriverGetVersion())
-    return _py_major_ver, _driver_ver
 
 
 @pytest.fixture(scope="function")
@@ -85,9 +71,7 @@ def test_get_kernel(init_cuda):
         ("cluster_scheduling_policy_preference", int),
     ],
 )
-def test_read_only_kernel_attributes(get_saxpy_kernel, attr, expected_type, cuda_version):
-    if cuda_version[0] < 12 and cuda_version[1] >= 12000:
-        pytest.skip("CUDA version is less than 12, and doesn't support kernel attribute access")
+def test_read_only_kernel_attributes(get_saxpy_kernel, attr, expected_type):
     kernel = get_saxpy_kernel
     method = getattr(kernel.attributes, attr)
     # get the value without providing a device ordinal