Implement Kernel.num_arguments, and Kernel.arguments_info

Author: oleksandr-pavlyk

cuda_core/cuda/core/experimental/_module.py

@@ -43,6 +43,7 @@ def _lazy_init():
             "data": driver.cuLibraryLoadData,
             "kernel": driver.cuLibraryGetKernel,
             "attribute": driver.cuKernelGetAttribute,
+            "paraminfo": driver.cuKernelGetParamInfo,
         }
         _kernel_ctypes = (driver.CUfunction, driver.CUkernel)
     else:
@@ -215,6 +216,36 @@ def attributes(self) -> KernelAttributes:
             self._attributes = KernelAttributes._init(self._handle)
         return self._attributes
 
+    @property
+    def num_arguments(self) -> int:
+        """int : The number of arguments of this function"""
+        attr_impl = self.attributes
+        if attr_impl._backend_version != "new":
+            raise NotImplementedError("New backend is required")
+        arg_pos = 0
+        while True:
+            result = attr_impl._loader["paraminfo"](self._handle, arg_pos)
+            if result[0] != driver.CUresult.CUDA_SUCCESS:
+                break
+            arg_pos = arg_pos + 1
+        return arg_pos
+
+    @property
+    def arguments_info(self) -> list[tuple[int, int]]:
+        """tuple[tuple[int, int]]: (offset, size) for each argument of this function"""
+        attr_impl = self.attributes
+        if attr_impl._backend_version != "new":
+            raise NotImplementedError("New backend is required")
+        arg_pos = 0
+        param_info = []
+        while True:
+            result = attr_impl._loader["paraminfo"](self._handle, arg_pos)
+            if result[0] != driver.CUresult.CUDA_SUCCESS:
+                break
+            param_info.append((result[1], result[2]))
+            arg_pos = arg_pos + 1
+        return param_info
+
     # TODO: implement from_handle()
 
 

cuda_core/tests/test_module.py

@@ -14,7 +14,7 @@
 import cuda.core.experimental
 from cuda.core.experimental import ObjectCode, Program, ProgramOptions, system
 
-SAXPY_KERNEL = """
+SAXPY_KERNEL = r"""
 template<typename T>
 __global__ void saxpy(const T a,
                     const T* x,
@@ -162,3 +162,50 @@ 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
+
+
+def test_saxpy_arguments(get_saxpy_kernel):
+    import ctypes
+    krn, _ = get_saxpy_kernel
+
+    assert krn.num_arguments == 5
+
+    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, sizes = zip(*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))
+
+
+@pytest.mark.parametrize("nargs", [0, 1, 2, 3, 16])
+def test_num_arguments(init_cuda, nargs):
+    args_str = ", ".join([f"int 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
+
+    import ctypes
+    class ExpectedStruct(ctypes.Structure):
+        _fields_ = [
+            (f'arg_{i}', ctypes.c_int) for i in range(nargs)
+        ]
+    members = tuple(getattr(ExpectedStruct, f"arg_{i}") for i in range(nargs))
+    expected = [(m.offset, m.size) for m in members]
+    assert krn.arguments_info == expected