numba-cuda 0.19.0__py3-none-any.whl → 0.20.0__py3-none-any.whl

numba_cuda/numba/cuda/core/funcdesc.py

@@ -0,0 +1,330 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+"""
+Function descriptors.
+"""
+
+from collections import defaultdict
+import importlib
+
+from numba.core import types
+from numba.cuda import itanium_mangler
+from numba.cuda.utils import _dynamic_modname, _dynamic_module
+
+
+def default_mangler(name, argtypes, *, abi_tags=(), uid=None):
+    return itanium_mangler.mangle(name, argtypes, abi_tags=abi_tags, uid=uid)
+
+
+def qualifying_prefix(modname, qualname):
+    """
+    Returns a new string that is used for the first half of the mangled name.
+    """
+    # XXX choose a different convention for object mode
+    return "{}.{}".format(modname, qualname) if modname else qualname
+
+
+class FunctionDescriptor(object):
+    """
+    Base class for function descriptors: an object used to carry
+    useful metadata about a natively callable function.
+
+    Note that while `FunctionIdentity` denotes a Python function
+    which is being concretely compiled by Numba, `FunctionDescriptor`
+    may be more "abstract".
+    """
+
+    __slots__ = (
+        "native",
+        "modname",
+        "qualname",
+        "doc",
+        "typemap",
+        "calltypes",
+        "args",
+        "kws",
+        "restype",
+        "argtypes",
+        "mangled_name",
+        "unique_name",
+        "env_name",
+        "global_dict",
+        "inline",
+        "noalias",
+        "abi_tags",
+        "uid",
+    )
+
+    def __init__(
+        self,
+        native,
+        modname,
+        qualname,
+        unique_name,
+        doc,
+        typemap,
+        restype,
+        calltypes,
+        args,
+        kws,
+        mangler=None,
+        argtypes=None,
+        inline=False,
+        noalias=False,
+        env_name=None,
+        global_dict=None,
+        abi_tags=(),
+        uid=None,
+    ):
+        self.native = native
+        self.modname = modname
+        self.global_dict = global_dict
+        self.qualname = qualname
+        self.unique_name = unique_name
+        self.doc = doc
+        # XXX typemap and calltypes should be on the compile result,
+        # not the FunctionDescriptor
+        self.typemap = typemap
+        self.calltypes = calltypes
+        self.args = args
+        self.kws = kws
+        self.restype = restype
+        # Argument types
+        if argtypes is not None:
+            assert isinstance(argtypes, tuple), argtypes
+            self.argtypes = argtypes
+        else:
+            # Get argument types from the type inference result
+            # (note the "arg.FOO" convention as used in typeinfer
+            self.argtypes = tuple(self.typemap["arg." + a] for a in args)
+        mangler = default_mangler if mangler is None else mangler
+        # The mangled name *must* be unique, else the wrong function can
+        # be chosen at link time.
+        qualprefix = qualifying_prefix(self.modname, self.qualname)
+        self.uid = uid
+        self.mangled_name = mangler(
+            qualprefix,
+            self.argtypes,
+            abi_tags=abi_tags,
+            uid=uid,
+        )
+        if env_name is None:
+            env_name = mangler(
+                ".NumbaEnv.{}".format(qualprefix),
+                self.argtypes,
+                abi_tags=abi_tags,
+                uid=uid,
+            )
+        self.env_name = env_name
+        self.inline = inline
+        self.noalias = noalias
+        self.abi_tags = abi_tags
+
+    def lookup_globals(self):
+        """
+        Return the global dictionary of the function.
+        It may not match the Module's globals if the function is created
+        dynamically (i.e. exec)
+        """
+        return self.global_dict or self.lookup_module().__dict__
+
+    def lookup_module(self):
+        """
+        Return the module in which this function is supposed to exist.
+        This may be a dummy module if the function was dynamically
+        generated or the module can't be found.
+        """
+        if self.modname == _dynamic_modname:
+            return _dynamic_module
+        else:
+            try:
+                # ensure module exist
+                return importlib.import_module(self.modname)
+            except ImportError:
+                return _dynamic_module
+
+    def lookup_function(self):
+        """
+        Return the original function object described by this object.
+        """
+        return getattr(self.lookup_module(), self.qualname)
+
+    @property
+    def llvm_func_name(self):
+        """
+        The LLVM-registered name for the raw function.
+        """
+        return self.mangled_name
+
+    @property
+    def llvm_cpython_wrapper_name(self):
+        """
+        The LLVM-registered name for a CPython-compatible wrapper of the
+        raw function (i.e. a PyCFunctionWithKeywords).
+        """
+        return itanium_mangler.prepend_namespace(
+            self.mangled_name, ns="cpython"
+        )
+
+    @property
+    def llvm_cfunc_wrapper_name(self):
+        """
+        The LLVM-registered name for a C-compatible wrapper of the
+        raw function.
+        """
+        return "cfunc." + self.mangled_name
+
+    def __repr__(self):
+        return "<function descriptor %r>" % (self.unique_name)
+
+    @classmethod
+    def _get_function_info(cls, func_ir):
+        """
+        Returns
+        -------
+        qualname, unique_name, modname, doc, args, kws, globals
+
+        ``unique_name`` must be a unique name.
+        """
+        func = func_ir.func_id.func
+        qualname = func_ir.func_id.func_qualname
+        # XXX to func_id
+        modname = func.__module__
+        doc = func.__doc__ or ""
+        args = tuple(func_ir.arg_names)
+        kws = ()  # TODO
+        global_dict = None
+
+        if modname is None:
+            # Dynamically generated function.
+            modname = _dynamic_modname
+            # Retain a reference to the dictionary of the function.
+            # This disables caching, serialization and pickling.
+            global_dict = func_ir.func_id.func.__globals__
+
+        unique_name = func_ir.func_id.unique_name
+
+        return qualname, unique_name, modname, doc, args, kws, global_dict
+
+    @classmethod
+    def _from_python_function(
+        cls,
+        func_ir,
+        typemap,
+        restype,
+        calltypes,
+        native,
+        mangler=None,
+        inline=False,
+        noalias=False,
+        abi_tags=(),
+    ):
+        (
+            qualname,
+            unique_name,
+            modname,
+            doc,
+            args,
+            kws,
+            global_dict,
+        ) = cls._get_function_info(func_ir)
+
+        self = cls(
+            native,
+            modname,
+            qualname,
+            unique_name,
+            doc,
+            typemap,
+            restype,
+            calltypes,
+            args,
+            kws,
+            mangler=mangler,
+            inline=inline,
+            noalias=noalias,
+            global_dict=global_dict,
+            abi_tags=abi_tags,
+            uid=func_ir.func_id.unique_id,
+        )
+        return self
+
+
+class PythonFunctionDescriptor(FunctionDescriptor):
+    """
+    A FunctionDescriptor subclass for Numba-compiled functions.
+    """
+
+    __slots__ = ()
+
+    @classmethod
+    def from_specialized_function(
+        cls,
+        func_ir,
+        typemap,
+        restype,
+        calltypes,
+        mangler,
+        inline,
+        noalias,
+        abi_tags,
+    ):
+        """
+        Build a FunctionDescriptor for a given specialization of a Python
+        function (in nopython mode).
+        """
+        return cls._from_python_function(
+            func_ir,
+            typemap,
+            restype,
+            calltypes,
+            native=True,
+            mangler=mangler,
+            inline=inline,
+            noalias=noalias,
+            abi_tags=abi_tags,
+        )
+
+    @classmethod
+    def from_object_mode_function(cls, func_ir):
+        """
+        Build a FunctionDescriptor for an object mode variant of a Python
+        function.
+        """
+        typemap = defaultdict(lambda: types.pyobject)
+        calltypes = typemap.copy()
+        restype = types.pyobject
+        return cls._from_python_function(
+            func_ir, typemap, restype, calltypes, native=False
+        )
+
+
+class ExternalFunctionDescriptor(FunctionDescriptor):
+    """
+    A FunctionDescriptor subclass for opaque external functions
+    (e.g. raw C functions).
+    """
+
+    __slots__ = ()
+
+    def __init__(self, name, restype, argtypes):
+        args = ["arg%d" % i for i in range(len(argtypes))]
+
+        def mangler(a, x, abi_tags, uid=None):
+            return a
+
+        super(ExternalFunctionDescriptor, self).__init__(
+            native=True,
+            modname=None,
+            qualname=name,
+            unique_name=name,
+            doc="",
+            typemap=None,
+            restype=restype,
+            calltypes=None,
+            args=args,
+            kws=None,
+            mangler=mangler,
+            argtypes=argtypes,
+        )