Implement Kernel.num_arguments, and Kernel.arguments_info (#612)

* Implement Kernel.num_arguments, and Kernel.arguments_info

* Factor out common logic between num_arguments and arguments_info properties

Also parametrize test to check arguments_info to check with all int,
and all short arguments.

* Use namedtuple turn to pack offset/size for individual kernel argument

* Add feature entry to 0.3.0 release notes

* Move ParamInfo to class context

Used modern namedtuple instance constructor.

* _get_arguments_info mustn't treat all error code as equal

As `arg_pos` is being incremented, eventually `CUDA_ERROR_INVALID_VALUE`
is expected when `arg_pos` exceeds the number of kernel arguments.

But other errors like CUDA_ERROR_INVALID_CONTEXT are possible, and they
should be treated as errors.

Failure to do that was behind the unintuitive behavior I reported in the PR.

* Add skip if testing with compute sanitizer, skip if cuda version < 12

Also add a test that if cuda is not initialized, an error is raised when
checking for the num_arguments.

* Implement suggestions from @leofang

1. Narrow fixture scope
2. Rename fixture to cuda12_prerequisite_check that provide a boolean
   rather than a pair of versions

* Use ParamInfo instead of NamedTuple is annotation

This required moving ParamInfo definition from class scope
to module scope, since referencing Kernel.ParamInfo from
annotations of methods of the Kernel class results in error
that Kernel class does not yet exist.

Author: oleksandr-pavlyk

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

@@ -26,24 +26,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

@@ -1,14 +1,17 @@
 # Copyright 2024 NVIDIA Corporation.  All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+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,
@@ -23,6 +26,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.
@@ -156,3 +168,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