Implement Kernel.num_arguments, and Kernel.arguments_info

cuda_core/cuda/core/experimental/_module.py

@@ -2,6 +2,7 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
+from collections import namedtuple
 from typing import Optional, Union
 from warnings import warn
 
@@ -43,6 +44,7 @@ def _lazy_init():
  "data": driver.cuLibraryLoadData,
  "kernel": driver.cuLibraryGetKernel,
  "attribute": driver.cuKernelGetAttribute,
+ "paraminfo": driver.cuKernelGetParamInfo,
  }
  _kernel_ctypes = (driver.CUfunction, driver.CUkernel)
  else:
@@ -182,6 +184,9 @@ def cluster_scheduling_policy_preference(self, device_id: int = None) -> int:
  )
 
 
+ParamInfo = namedtuple("ParamInfo", ["offset", "size"])
+
+
 class Kernel:
  """Represent a compiled kernel that had been loaded onto the device.
 
@@ -215,6 +220,36 @@ def attributes(self) -> KernelAttributes:
  self._attributes = KernelAttributes._init(self._handle)
  return self._attributes
 
+ def _get_arguments_info(self, param_info=False) -> tuple[int, list[ParamInfo]]:
+ attr_impl = self.attributes
+ if attr_impl._backend_version != "new":
+ raise NotImplementedError("New backend is required")
+ arg_pos = 0
+ param_info_data = []
+ while True:
+ result = attr_impl._loader["paraminfo"](self._handle, arg_pos)
+ if result[0] != driver.CUresult.CUDA_SUCCESS:
+ break
+ if param_info:
+ p_info = ParamInfo(offset=result[1], size=result[2])
+ param_info_data.append(p_info)
+ arg_pos = arg_pos + 1
+ if result[0] != driver.CUresult.CUDA_ERROR_INVALID_VALUE:
+ handle_return(result)
+ return arg_pos, param_info_data
+
+ @property
+ def num_arguments(self) -> int:
+ """int : The number of arguments of this function"""
+ num_args, _ = self._get_arguments_info()
+ return num_args
+
+ @property
+ def arguments_info(self) -> list[ParamInfo]:
+ """list[ParamInfo]: (offset, size) for each argument of this function"""
+ _, param_info = self._get_arguments_info(param_info=True)
+ return param_info
+
  # TODO: implement from_handle()
 
 

cuda_core/docs/source/release/0.3.0-notes.rst

@@ -20,6 +20,7 @@ Breaking Changes
 New features
 ------------
 
+- :class:`Kernel` adds :property:`Kernel.num_arguments` and :property:`Kernel.arguments_info` for introspection of kernel arguments. (#612)
 
 New examples
 ------------

cuda_core/tests/conftest.py

@@ -31,24 +31,45 @@ def init_cuda():
  device = Device()
  device.set_current()
  yield
- _device_unset_current()
+_ = _device_unset_current()
 
 
-def _device_unset_current():
+def _device_unset_current() -> bool:
+ """Pop current CUDA context.
+
+ Returns True if context was popped, False it the stack was empty.
+ """
  ctx = handle_return(driver.cuCtxGetCurrent())
  if int(ctx) == 0:
  # no active context, do nothing
- return
+ return False
  handle_return(driver.cuCtxPopCurrent())
  if hasattr(_device._tls, "devices"):
  del _device._tls.devices
+ return True
 
 
 @pytest.fixture(scope="function")
 def deinit_cuda():
  # TODO: rename this to e.g. deinit_context
  yield
- _device_unset_current()
+ _ = _device_unset_current()
+
+
+@pytest.fixture(scope="function")
+def deinit_all_contexts_function():
+ def pop_all_contexts():
+ max_iters = 256
+ for _ in range(max_iters):
+ if _device_unset_current():
+ # context was popped, continue until stack is empty
+ continue
+ # no active context, we are ready
+ break
+ else:
+ raise RuntimeError(f"Number of iterations popping current CUDA contexts, exceded {max_iters}")
+
+ return pop_all_contexts
 
 
 # samples relying on cffi could fail as the modules cannot be imported

cuda_core/tests/test_module.py

@@ -7,14 +7,17 @@
 # is strictly prohibited.
 
 
+import ctypes
 import warnings
 
 import pytest
+from conftest import skipif_testing_with_compute_sanitizer
 
 import cuda.core.experimental
 from cuda.core.experimental import ObjectCode, Program, ProgramOptions, system
+from cuda.core.experimental._utils.cuda_utils import CUDAError, driver, get_binding_version, handle_return
 
-SAXPY_KERNEL = """
+SAXPY_KERNEL = r"""
 template<typename T>
 __global__ void saxpy(const T a,
  const T* x,
@@ -29,6 +32,15 @@
 """
 
 
+@pytest.fixture(scope="module")
+def cuda12_prerequisite_check():
+ # binding availability depends on cuda-python version
+ # and version of underlying CUDA toolkit
+ _py_major_ver, _ = get_binding_version()
+ _driver_ver = handle_return(driver.cuDriverGetVersion())
+ return _py_major_ver >= 12 and _driver_ver >= 12000
+
+
 def test_kernel_attributes_init_disabled():
  with pytest.raises(RuntimeError, match=r"^KernelAttributes cannot be instantiated directly\."):
  cuda.core.experimental._module.KernelAttributes() # Ensure back door is locked.
@@ -162,3 +174,80 @@ def test_object_code_load_cubin_from_file(get_saxpy_kernel, tmp_path):
 def test_object_code_handle(get_saxpy_object_code):
  mod = get_saxpy_object_code
  assert mod.handle is not None
+
+
+@skipif_testing_with_compute_sanitizer
+def test_saxpy_arguments(get_saxpy_kernel, cuda12_prerequisite_check):
+ if not cuda12_prerequisite_check:
+ pytest.skip("Test requires CUDA 12")
+ krn, _ = get_saxpy_kernel
+
+ assert krn.num_arguments == 5
+
+ assert "ParamInfo" in str(type(krn).arguments_info.fget.__annotations__)
+ arg_info = krn.arguments_info
+ n_args = len(arg_info)
+ assert n_args == krn.num_arguments
+
+ class ExpectedStruct(ctypes.Structure):
+ _fields_ = [
+ ("a", ctypes.c_float),
+ ("x", ctypes.POINTER(ctypes.c_float)),
+ ("y", ctypes.POINTER(ctypes.c_float)),
+ ("out", ctypes.POINTER(ctypes.c_float)),
+ ("N", ctypes.c_size_t),
+ ]
+
+ offsets = [p.offset for p in arg_info]
+ sizes = [p.size for p in arg_info]
+ members = [getattr(ExpectedStruct, name) for name, _ in ExpectedStruct._fields_]
+ expected_offsets = tuple(m.offset for m in members)
+ assert all(actual == expected for actual, expected in zip(offsets, expected_offsets))
+ expected_sizes = tuple(m.size for m in members)
+ assert all(actual == expected for actual, expected in zip(sizes, expected_sizes))
+
+
+@skipif_testing_with_compute_sanitizer
+@pytest.mark.parametrize("nargs", [0, 1, 2, 3, 16])
+@pytest.mark.parametrize("c_type_name,c_type", [("int", ctypes.c_int), ("short", ctypes.c_short)], ids=["int", "short"])
+def test_num_arguments(init_cuda, nargs, c_type_name, c_type, cuda12_prerequisite_check):
+ if not cuda12_prerequisite_check:
+ pytest.skip("Test requires CUDA 12")
+ args_str = ", ".join([f"{c_type_name} p_{i}" for i in range(nargs)])
+ src = f"__global__ void foo{nargs}({args_str}) {{ }}"
+ prog = Program(src, code_type="c++")
+ mod = prog.compile(
+ "cubin",
+ name_expressions=(f"foo{nargs}",),
+ )
+ krn = mod.get_kernel(f"foo{nargs}")
+ assert krn.num_arguments == nargs
+
+ class ExpectedStruct(ctypes.Structure):
+ _fields_ = [(f"arg_{i}", c_type) for i in range(nargs)]
+
+ members = tuple(getattr(ExpectedStruct, f"arg_{i}") for i in range(nargs))
+
+ arg_info = krn.arguments_info
+ assert all([actual.offset == expected.offset for actual, expected in zip(arg_info, members)])
+ assert all([actual.size == expected.size for actual, expected in zip(arg_info, members)])
+
+
+@skipif_testing_with_compute_sanitizer
+def test_num_args_error_handling(deinit_all_contexts_function, cuda12_prerequisite_check):
+ if not cuda12_prerequisite_check:
+ pytest.skip("Test requires CUDA 12")
+ src = "__global__ void foo(int a) { }"
+ prog = Program(src, code_type="c++")
+ mod = prog.compile(
+ "cubin",
+ name_expressions=("foo",),
+ )
+ krn = mod.get_kernel("foo")
+ # empty driver's context stack using function from conftest
+ deinit_all_contexts_function()
+ # with no current context, cuKernelGetParamInfo would report
+ # exception which we expect to handle by raising
+ with pytest.raises(CUDAError):
+ # assignment resolves linter error "B018: useless expression"
+ _ = krn.num_arguments