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

Author: oleksandr-pavlyk

cuda_core/cuda/core/experimental/_module.py

@@ -2,7 +2,8 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-from typing import Optional, Union
+from collections import namedtuple
+from typing import NamedTuple, Optional, Union
 from warnings import warn
 
 from cuda.core.experimental._utils.clear_error_support import (
@@ -216,18 +217,20 @@ def attributes(self) -> KernelAttributes:
             self._attributes = KernelAttributes._init(self._handle)
         return self._attributes
 
-    def _get_arguments_info(self, param_info=False) -> tuple[int, list[tuple[int, int]]]:
+    def _get_arguments_info(self, param_info=False) -> tuple[int, list[NamedTuple]]:
         attr_impl = self.attributes
         if attr_impl._backend_version != "new":
             raise NotImplementedError("New backend is required")
         arg_pos = 0
+        if param_info:
+            ParamInfo = namedtuple("ParamInfo", ["offset", "size"])
         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:
-                param_info_data.append(result[1:3])
+                param_info_data.append(ParamInfo._make(result[1:3]))
             arg_pos = arg_pos + 1
         return arg_pos, param_info_data
 
@@ -238,8 +241,8 @@ def num_arguments(self) -> int:
         return num_args
 
     @property
-    def arguments_info(self) -> list[tuple[int, int]]:
-        """list[tuple[int, int]]: (offset, size) for each argument of this function"""
+    def arguments_info(self) -> list[NamedTuple]:
+        """list[NamedTuple[int, int]]: (offset, size) for each argument of this function"""
         _, param_info = self._get_arguments_info(param_info=True)
         return param_info
 

cuda_core/tests/test_module.py

@@ -7,8 +7,8 @@
 # is strictly prohibited.
 
 
-import warnings
 import ctypes
+import warnings
 
 import pytest
 
@@ -173,15 +173,18 @@ def test_saxpy_arguments(get_saxpy_kernel):
     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)
+            ("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)
+
+    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))
@@ -197,15 +200,16 @@ def test_num_arguments(init_cuda, nargs, c_type_name, c_type):
     prog = Program(src, code_type="c++")
     mod = prog.compile(
         "cubin",
-        name_expressions=(f"foo{nargs}", ),
+        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)
-        ]
+        _fields_ = [(f"arg_{i}", c_type) 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
+
+    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)])