numba-cuda 0.22.0__cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl

numba_cuda/numba/cuda/datamodel/cuda_packer.py

@@ -0,0 +1,224 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+from collections import deque
+
+from numba.cuda import types
+from numba.cuda import cgutils
+
+
+class DataPacker(object):
+    """
+    A helper to pack a number of typed arguments into a data structure.
+    Omitted arguments (i.e. values with the type `Omitted`) are automatically
+    skipped.
+    """
+
+    # XXX should DataPacker be a model for a dedicated type?
+
+    def __init__(self, dmm, fe_types):
+        self._dmm = dmm
+        self._fe_types = fe_types
+        self._models = [dmm.lookup(ty) for ty in fe_types]
+
+        self._pack_map = []
+        self._be_types = []
+        for i, ty in enumerate(fe_types):
+            if not isinstance(ty, types.Omitted):
+                self._pack_map.append(i)
+                self._be_types.append(self._models[i].get_data_type())
+
+    def as_data(self, builder, values):
+        """
+        Return the given values packed as a data structure.
+        """
+        elems = [
+            self._models[i].as_data(builder, values[i]) for i in self._pack_map
+        ]
+        return cgutils.make_anonymous_struct(builder, elems)
+
+    def _do_load(self, builder, ptr, formal_list=None):
+        res = []
+        for i, i_formal in enumerate(self._pack_map):
+            elem_ptr = cgutils.gep_inbounds(builder, ptr, 0, i)
+            val = self._models[i_formal].load_from_data_pointer(
+                builder, elem_ptr
+            )
+            if formal_list is None:
+                res.append((self._fe_types[i_formal], val))
+            else:
+                formal_list[i_formal] = val
+        return res
+
+    def load(self, builder, ptr):
+        """
+        Load the packed values and return a (type, value) tuples.
+        """
+        return self._do_load(builder, ptr)
+
+    def load_into(self, builder, ptr, formal_list):
+        """
+        Load the packed values into a sequence indexed by formal
+        argument number (skipping any Omitted position).
+        """
+        self._do_load(builder, ptr, formal_list)
+
+
+class ArgPacker(object):
+    """
+    Compute the position for each high-level typed argument.
+    It flattens every composite argument into primitive types.
+    It maintains a position map for unflattening the arguments.
+
+    Since struct (esp. nested struct) have specific ABI requirements (e.g.
+    alignment, pointer address-space, ...) in different architecture (e.g.
+    OpenCL, CUDA), flattening composite argument types simplifes the call
+    setup from the Python side.  Functions are receiving simple primitive
+    types and there are only a handful of these.
+    """
+
+    def __init__(self, dmm, fe_args):
+        self._dmm = dmm
+        self._fe_args = fe_args
+        self._nargs = len(fe_args)
+
+        self._dm_args = []
+        argtys = []
+        for ty in fe_args:
+            dm = self._dmm.lookup(ty)
+            self._dm_args.append(dm)
+            argtys.append(dm.get_argument_type())
+        self._unflattener = _Unflattener(argtys)
+        self._be_args = list(_flatten(argtys))
+
+    def as_arguments(self, builder, values):
+        """Flatten all argument values"""
+        if len(values) != self._nargs:
+            raise TypeError(
+                "invalid number of args: expected %d, got %d"
+                % (self._nargs, len(values))
+            )
+
+        if not values:
+            return ()
+
+        args = [
+            dm.as_argument(builder, val)
+            for dm, val in zip(self._dm_args, values)
+        ]
+
+        args = tuple(_flatten(args))
+        return args
+
+    def from_arguments(self, builder, args):
+        """Unflatten all argument values"""
+
+        valtree = self._unflattener.unflatten(args)
+        values = [
+            dm.from_argument(builder, val)
+            for dm, val in zip(self._dm_args, valtree)
+        ]
+
+        return values
+
+    def assign_names(self, args, names):
+        """Assign names for each flattened argument values."""
+
+        valtree = self._unflattener.unflatten(args)
+        for aval, aname in zip(valtree, names):
+            self._assign_names(aval, aname)
+
+    def _assign_names(self, val_or_nested, name, depth=()):
+        if isinstance(val_or_nested, (tuple, list)):
+            for pos, aval in enumerate(val_or_nested):
+                self._assign_names(aval, name, depth=depth + (pos,))
+        else:
+            postfix = ".".join(map(str, depth))
+            parts = [name, postfix]
+            val_or_nested.name = ".".join(filter(bool, parts))
+
+    @property
+    def argument_types(self):
+        """Return a list of LLVM types that are results of flattening
+        composite types.
+        """
+        return tuple(ty for ty in self._be_args if ty != ())
+
+
+def _flatten(iterable):
+    """
+    Flatten nested iterable of (tuple, list).
+    """
+
+    def rec(iterable):
+        for i in iterable:
+            if isinstance(i, (tuple, list)):
+                for j in rec(i):
+                    yield j
+            else:
+                yield i
+
+    return rec(iterable)
+
+
+_PUSH_LIST = 1
+_APPEND_NEXT_VALUE = 2
+_APPEND_EMPTY_TUPLE = 3
+_POP = 4
+
+
+class _Unflattener(object):
+    """
+    An object used to unflatten nested sequences after a given pattern
+    (an arbitrarily nested sequence).
+    The pattern shows the nested sequence shape desired when unflattening;
+    the values it contains are irrelevant.
+    """
+
+    def __init__(self, pattern):
+        self._code = self._build_unflatten_code(pattern)
+
+    def _build_unflatten_code(self, iterable):
+        """Build the unflatten opcode sequence for the given *iterable* structure
+        (an iterable of nested sequences).
+        """
+        code = []
+
+        def rec(iterable):
+            for i in iterable:
+                if isinstance(i, (tuple, list)):
+                    if len(i) > 0:
+                        code.append(_PUSH_LIST)
+                        rec(i)
+                        code.append(_POP)
+                    else:
+                        code.append(_APPEND_EMPTY_TUPLE)
+                else:
+                    code.append(_APPEND_NEXT_VALUE)
+
+        rec(iterable)
+        return code
+
+    def unflatten(self, flatiter):
+        """Rebuild a nested tuple structure."""
+        vals = deque(flatiter)
+
+        res = []
+        cur = res
+        stack = []
+        for op in self._code:
+            if op is _PUSH_LIST:
+                stack.append(cur)
+                cur.append([])
+                cur = cur[-1]
+            elif op is _APPEND_NEXT_VALUE:
+                cur.append(vals.popleft())
+            elif op is _APPEND_EMPTY_TUPLE:
+                cur.append(())
+            elif op is _POP:
+                cur = stack.pop()
+
+        assert not stack, stack
+        assert not vals, vals
+
+        return res

numba_cuda/numba/cuda/datamodel/cuda_registry.py

@@ -0,0 +1,22 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import functools
+from .manager import DataModelManager
+
+
+def register(dmm, typecls):
+    """Used as decorator to simplify datamodel registration.
+    Returns the object being decorated so that chaining is possible.
+    """
+
+    def wraps(fn):
+        dmm.register(typecls, fn)
+        return fn
+
+    return wraps
+
+
+default_manager = DataModelManager()
+
+register_default = functools.partial(register, default_manager)

numba_cuda/numba/cuda/datamodel/cuda_testing.py

@@ -0,0 +1,153 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+from llvmlite import ir
+from llvmlite import binding as ll
+
+from numba.cuda import datamodel
+import unittest
+
+
+class DataModelTester(unittest.TestCase):
+    """
+    Test the implementation of a DataModel for a frontend type.
+    """
+
+    fe_type = NotImplemented
+
+    def setUp(self):
+        self.module = ir.Module()
+        self.datamodel = datamodel.default_manager[self.fe_type]
+
+    def test_as_arg(self):
+        """
+        - Is as_arg() and from_arg() implemented?
+        - Are they the inverse of each other?
+        """
+        fnty = ir.FunctionType(ir.VoidType(), [])
+        function = ir.Function(self.module, fnty, name="test_as_arg")
+        builder = ir.IRBuilder()
+        builder.position_at_end(function.append_basic_block())
+
+        undef_value = ir.Constant(self.datamodel.get_value_type(), None)
+        args = self.datamodel.as_argument(builder, undef_value)
+        self.assertIsNot(
+            args, NotImplemented, "as_argument returned NotImplementedError"
+        )
+
+        if isinstance(args, (tuple, list)):
+
+            def recur_tuplize(args, func=None):
+                for arg in args:
+                    if isinstance(arg, (tuple, list)):
+                        yield tuple(recur_tuplize(arg, func=func))
+                    else:
+                        if func is None:
+                            yield arg
+                        else:
+                            yield func(arg)
+
+            argtypes = tuple(recur_tuplize(args, func=lambda x: x.type))
+            exptypes = tuple(recur_tuplize(self.datamodel.get_argument_type()))
+            self.assertEqual(exptypes, argtypes)
+        else:
+            self.assertEqual(args.type, self.datamodel.get_argument_type())
+
+        rev_value = self.datamodel.from_argument(builder, args)
+        self.assertEqual(rev_value.type, self.datamodel.get_value_type())
+
+        builder.ret_void()  # end function
+
+        # Ensure valid LLVM generation
+        materialized = ll.parse_assembly(str(self.module))
+        str(materialized)
+
+    def test_as_return(self):
+        """
+        - Is as_return() and from_return() implemented?
+        - Are they the inverse of each other?
+        """
+        fnty = ir.FunctionType(ir.VoidType(), [])
+        function = ir.Function(self.module, fnty, name="test_as_return")
+        builder = ir.IRBuilder()
+        builder.position_at_end(function.append_basic_block())
+
+        undef_value = ir.Constant(self.datamodel.get_value_type(), None)
+        ret = self.datamodel.as_return(builder, undef_value)
+        self.assertIsNot(
+            ret, NotImplemented, "as_return returned NotImplementedError"
+        )
+
+        self.assertEqual(ret.type, self.datamodel.get_return_type())
+
+        rev_value = self.datamodel.from_return(builder, ret)
+        self.assertEqual(rev_value.type, self.datamodel.get_value_type())
+
+        builder.ret_void()  # end function
+
+        # Ensure valid LLVM generation
+        materialized = ll.parse_assembly(str(self.module))
+        str(materialized)
+
+
+class SupportAsDataMixin(object):
+    """Test as_data() and from_data()"""
+
+    # XXX test load_from_data_pointer() as well
+
+    def test_as_data(self):
+        fnty = ir.FunctionType(ir.VoidType(), [])
+        function = ir.Function(self.module, fnty, name="test_as_data")
+        builder = ir.IRBuilder()
+        builder.position_at_end(function.append_basic_block())
+
+        undef_value = ir.Constant(self.datamodel.get_value_type(), None)
+        data = self.datamodel.as_data(builder, undef_value)
+        self.assertIsNot(
+            data, NotImplemented, "as_data returned NotImplemented"
+        )
+
+        self.assertEqual(data.type, self.datamodel.get_data_type())
+
+        rev_value = self.datamodel.from_data(builder, data)
+        self.assertEqual(rev_value.type, self.datamodel.get_value_type())
+
+        builder.ret_void()  # end function
+
+        # Ensure valid LLVM generation
+        materialized = ll.parse_assembly(str(self.module))
+        str(materialized)
+
+
+class NotSupportAsDataMixin(object):
+    """Ensure as_data() and from_data() raise NotImplementedError."""
+
+    def test_as_data_not_supported(self):
+        fnty = ir.FunctionType(ir.VoidType(), [])
+        function = ir.Function(self.module, fnty, name="test_as_data")
+        builder = ir.IRBuilder()
+        builder.position_at_end(function.append_basic_block())
+
+        undef_value = ir.Constant(self.datamodel.get_value_type(), None)
+        with self.assertRaises(NotImplementedError):
+            data = self.datamodel.as_data(builder, undef_value)  # noqa: F841
+        with self.assertRaises(NotImplementedError):
+            rev_data = self.datamodel.from_data(builder, undef_value)  # noqa: F841
+
+
+class DataModelTester_SupportAsDataMixin(DataModelTester, SupportAsDataMixin):
+    pass
+
+
+class DataModelTester_NotSupportAsDataMixin(
+    DataModelTester, NotSupportAsDataMixin
+):
+    pass
+
+
+def test_factory(support_as_data=True):
+    """A helper for returning a unittest TestCase for testing"""
+    if support_as_data:
+        return DataModelTester_SupportAsDataMixin
+    else:
+        return DataModelTester_NotSupportAsDataMixin

numba_cuda/numba/cuda/datamodel/manager.py

@@ -0,0 +1,11 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import sys
+from numba.cuda.utils import redirect_numba_module
+
+sys.modules[__name__] = redirect_numba_module(
+    locals(),
+    "numba.core.datamodel.manager",
+    "numba.cuda.datamodel.cuda_manager",
+)

numba_cuda/numba/cuda/datamodel/models.py

@@ -0,0 +1,9 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import sys
+from numba.cuda.utils import redirect_numba_module
+
+sys.modules[__name__] = redirect_numba_module(
+    locals(), "numba.core.datamodel.models", "numba.cuda.datamodel.cuda_models"
+)

numba_cuda/numba/cuda/datamodel/packer.py

@@ -0,0 +1,9 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import sys
+from numba.cuda.utils import redirect_numba_module
+
+sys.modules[__name__] = redirect_numba_module(
+    locals(), "numba.core.datamodel.packer", "numba.cuda.datamodel.cuda_packer"
+)

numba_cuda/numba/cuda/datamodel/registry.py

@@ -0,0 +1,11 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import sys
+from numba.cuda.utils import redirect_numba_module
+
+sys.modules[__name__] = redirect_numba_module(
+    locals(),
+    "numba.core.datamodel.registry",
+    "numba.cuda.datamodel.cuda_registry",
+)

numba_cuda/numba/cuda/datamodel/testing.py

@@ -0,0 +1,11 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import sys
+from numba.cuda.utils import redirect_numba_module
+
+sys.modules[__name__] = redirect_numba_module(
+    locals(),
+    "numba.core.datamodel.testing",
+    "numba.cuda.datamodel.cuda_testing",
+)