numba-cuda 0.22.0__cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl

numba_cuda/numba/cuda/core/targetconfig.py

@@ -0,0 +1,330 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+"""
+This module contains utils for manipulating target configurations such as
+compiler flags.
+"""
+
+import re
+import zlib
+import base64
+
+from types import MappingProxyType
+from numba.cuda import utils
+from numba.cuda import HAS_NUMBA
+
+
+class Option:
+    """An option to be used in ``TargetConfig``."""
+
+    __slots__ = "_type", "_default", "_doc"
+
+    def __init__(self, type, *, default, doc):
+        """
+        Parameters
+        ----------
+        type :
+            Type of the option value. It can be a callable.
+            The setter always calls ``self._type(value)``.
+        default :
+            The default value for the option.
+        doc : str
+            Docstring for the option.
+        """
+        self._type = type
+        self._default = default
+        self._doc = doc
+
+    @property
+    def type(self):
+        return self._type
+
+    @property
+    def default(self):
+        return self._default
+
+    @property
+    def doc(self):
+        return self._doc
+
+
+if HAS_NUMBA:
+    from numba.core.targetconfig import ConfigStack, _FlagsStack
+else:
+
+    class _FlagsStack(utils.ThreadLocalStack, stack_name="flags"):
+        pass
+
+    class ConfigStack:
+        """A stack for tracking target configurations in the compiler.
+
+        It stores the stack in a thread-local class attribute. All instances in the
+        same thread will see the same stack.
+        """
+
+        @classmethod
+        def top_or_none(cls):
+            """Get the TOS or return None if no config is set."""
+            self = cls()
+            if self:
+                flags = self.top()
+            else:
+                # Note: should this be the default flag for the target instead?
+                flags = None
+            return flags
+
+        def __init__(self):
+            self._stk = _FlagsStack()
+
+        def top(self):
+            return self._stk.top()
+
+        def __len__(self):
+            return len(self._stk)
+
+        def enter(self, flags):
+            """Returns a contextmanager that performs ``push(flags)`` on enter and
+            ``pop()`` on exit.
+            """
+            return self._stk.enter(flags)
+
+
+class _MetaTargetConfig(type):
+    """Metaclass for ``TargetConfig``.
+
+    When a subclass of ``TargetConfig`` is created, all ``Option`` defined
+    as class members will be parsed and corresponding getters, setters, and
+    delters will be inserted.
+    """
+
+    def __init__(cls, name, bases, dct):
+        """Invoked when subclass is created.
+
+        Insert properties for each ``Option`` that are class members.
+        All the options will be grouped inside the ``.options`` class
+        attribute.
+        """
+        # Gather options from base classes and class dict
+        opts = {}
+        # Reversed scan into the base classes to follow MRO ordering such that
+        # the closest base class is overriding
+        for base_cls in reversed(bases):
+            opts.update(base_cls.options)
+        opts.update(cls.find_options(dct))
+        # Store the options into class attribute as a ready-only mapping.
+        cls.options = MappingProxyType(opts)
+
+        # Make properties for each of the options
+        def make_prop(name, option):
+            def getter(self):
+                return self._values.get(name, option.default)
+
+            def setter(self, val):
+                self._values[name] = option.type(val)
+
+            def delter(self):
+                del self._values[name]
+
+            return property(getter, setter, delter, option.doc)
+
+        for name, option in cls.options.items():
+            setattr(cls, name, make_prop(name, option))
+
+    def find_options(cls, dct):
+        """Returns a new dict with all the items that are a mapping to an
+        ``Option``.
+        """
+        return {k: v for k, v in dct.items() if isinstance(v, Option)}
+
+
+class _NotSetType:
+    def __repr__(self):
+        return "<NotSet>"
+
+
+_NotSet = _NotSetType()
+
+
+class TargetConfig(metaclass=_MetaTargetConfig):
+    """Base class for ``TargetConfig``.
+
+    Subclass should fill class members with ``Option``. For example:
+
+    >>> class MyTargetConfig(TargetConfig):
+    >>>     a_bool_option = Option(type=bool, default=False, doc="a bool")
+    >>>     an_int_option = Option(type=int, default=0, doc="an int")
+
+    The metaclass will insert properties for each ``Option``. For example:
+
+    >>> tc = MyTargetConfig()
+    >>> tc.a_bool_option = True  # invokes the setter
+    >>> print(tc.an_int_option)  # print the default
+    """
+
+    __slots__ = ["_values"]
+
+    # Used for compression in mangling.
+    # Set to -15 to disable the header and checksum for smallest output.
+    _ZLIB_CONFIG = {"wbits": -15}
+
+    def __init__(self, copy_from=None):
+        """
+        Parameters
+        ----------
+        copy_from : TargetConfig or None
+            if None, creates an empty ``TargetConfig``.
+            Otherwise, creates a copy.
+        """
+        self._values = {}
+        if copy_from is not None:
+            assert isinstance(copy_from, TargetConfig)
+            self._values.update(copy_from._values)
+
+    def __repr__(self):
+        # NOTE: default options will be placed at the end and grouped inside
+        #       a square bracket; i.e. [optname=optval, ...]
+        args = []
+        defs = []
+        for k in self.options:
+            msg = f"{k}={getattr(self, k)}"
+            if not self.is_set(k):
+                defs.append(msg)
+            else:
+                args.append(msg)
+        clsname = self.__class__.__name__
+        return f"{clsname}({', '.join(args)}, [{', '.join(defs)}])"
+
+    def __hash__(self):
+        return hash(tuple(sorted(self.values())))
+
+    def __eq__(self, other):
+        if isinstance(other, TargetConfig):
+            return self.values() == other.values()
+        else:
+            return NotImplemented
+
+    def values(self):
+        """Returns a dict of all the values"""
+        return {k: getattr(self, k) for k in self.options}
+
+    def is_set(self, name):
+        """Is the option set?"""
+        self._guard_option(name)
+        return name in self._values
+
+    def discard(self, name):
+        """Remove the option by name if it is defined.
+
+        After this, the value for the option will be set to its default value.
+        """
+        self._guard_option(name)
+        self._values.pop(name, None)
+
+    def inherit_if_not_set(self, name, default=_NotSet):
+        """Inherit flag from ``ConfigStack``.
+
+        Parameters
+        ----------
+        name : str
+            Option name.
+        default : optional
+            When given, it overrides the default value.
+            It is only used when the flag is not defined locally and there is
+            no entry in the ``ConfigStack``.
+        """
+        self._guard_option(name)
+        if not self.is_set(name):
+            cstk = ConfigStack()
+            if cstk:
+                # inherit
+                top = cstk.top()
+                setattr(self, name, getattr(top, name))
+            elif default is not _NotSet:
+                setattr(self, name, default)
+
+    def copy(self):
+        """Clone this instance."""
+        return type(self)(self)
+
+    def summary(self) -> str:
+        """Returns a ``str`` that summarizes this instance.
+
+        In contrast to ``__repr__``, only options that are explicitly set will
+        be shown.
+        """
+        args = [f"{k}={v}" for k, v in self._summary_args()]
+        clsname = self.__class__.__name__
+        return f"{clsname}({', '.join(args)})"
+
+    def _guard_option(self, name):
+        if name not in self.options:
+            msg = f"{name!r} is not a valid option for {type(self)}"
+            raise ValueError(msg)
+
+    def _summary_args(self):
+        """returns a sorted sequence of 2-tuple containing the
+        ``(flag_name, flag_value)`` for flag that are set with a non-default
+        value.
+        """
+        args = []
+        for k in sorted(self.options):
+            opt = self.options[k]
+            if self.is_set(k):
+                flagval = getattr(self, k)
+                if opt.default != flagval:
+                    v = (k, flagval)
+                    args.append(v)
+        return args
+
+    @classmethod
+    def _make_compression_dictionary(cls) -> bytes:
+        """Returns a ``bytes`` object suitable for use as a dictionary for
+        compression.
+        """
+        buf = []
+        # include package name
+        buf.append("numba")
+        # include class name
+        buf.append(cls.__class__.__name__)
+        # include common values
+        buf.extend(["True", "False"])
+        # include all options name and their default value
+        for k, opt in cls.options.items():
+            buf.append(k)
+            buf.append(str(opt.default))
+        return "".join(buf).encode()
+
+    def get_mangle_string(self) -> str:
+        """Return a string suitable for symbol mangling."""
+        zdict = self._make_compression_dictionary()
+
+        comp = zlib.compressobj(
+            zdict=zdict, level=zlib.Z_BEST_COMPRESSION, **self._ZLIB_CONFIG
+        )
+        # The mangled string is a compressed and base64 encoded version of the
+        # summary
+        buf = [comp.compress(self.summary().encode())]
+        buf.append(comp.flush())
+        return base64.b64encode(b"".join(buf)).decode()
+
+    @classmethod
+    def demangle(cls, mangled: str) -> str:
+        """Returns the demangled result from ``.get_mangle_string()``"""
+
+        # unescape _XX sequence
+        def repl(x):
+            return chr(int("0x" + x.group(0)[1:], 16))
+
+        unescaped = re.sub(r"_[a-zA-Z0-9][a-zA-Z0-9]", repl, mangled)
+        # decode base64
+        raw = base64.b64decode(unescaped)
+        # decompress
+        zdict = cls._make_compression_dictionary()
+        dc = zlib.decompressobj(zdict=zdict, **cls._ZLIB_CONFIG)
+        buf = []
+        while raw:
+            buf.append(dc.decompress(raw))
+            raw = dc.unconsumed_tail
+        buf.append(dc.flush())
+        return b"".join(buf).decode()

numba_cuda/numba/cuda/core/tracing.py

@@ -0,0 +1,231 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import inspect
+import logging
+import sys
+import threading
+from functools import wraps
+from itertools import chain
+
+from numba.cuda.core import config
+
+
+class TLS(threading.local):
+    """Use a subclass to properly initialize the TLS variables in all threads."""  # noqa: E501
+
+    def __init__(self):
+        self.tracing = False
+        self.indent = 0
+
+
+tls = TLS()
+
+
+def find_function_info(func, spec, args):
+    """Return function meta-data in a tuple.
+
+    (name, type)"""
+
+    module = getattr(func, "__module__", None)
+    name = getattr(func, "__name__", None)
+    self = getattr(func, "__self__", None)
+    cname = None
+    if self:
+        cname = self.__name__
+        # cname = self.__class__.__name__
+    # Try to deduce the class' name even for unbound methods from their
+    # first argument, which we assume to be a class instance if named 'self'...
+    elif len(spec.args) and spec.args[0] == "self":
+        cname = args[0].__class__.__name__
+    # ...or a class object if named 'cls'
+    elif len(spec.args) and spec.args[0] == "cls":
+        cname = args[0].__name__
+    if name:
+        qname = []
+        if module and module != "__main__":
+            qname.append(module)
+            qname.append(".")
+        if cname:
+            qname.append(cname)
+            qname.append(".")
+        qname.append(name)
+        name = "".join(qname)
+    return name, None
+
+
+def chop(value):
+    MAX_SIZE = 320
+    s = repr(value)
+    if len(s) > MAX_SIZE:
+        return s[:MAX_SIZE] + "..." + s[-1]
+    else:
+        return s
+
+
+def create_events(fname, spec, args, kwds):
+    values = dict()
+    if spec.defaults:
+        values = dict(zip(spec.args[-len(spec.defaults) :], spec.defaults))
+    values.update(kwds)
+    values.update(list(zip(spec.args[: len(args)], args)))
+    positional = ["%s=%r" % (a, values.pop(a)) for a in spec.args]
+    anonymous = [str(a) for a in args[len(positional) :]]
+    keywords = ["%s=%r" % (k, values[k]) for k in sorted(values.keys())]
+    params = ", ".join([f for f in chain(positional, anonymous, keywords) if f])
+
+    enter = [">> ", tls.indent * " ", fname, "(", params, ")"]
+    leave = ["<< ", tls.indent * " ", fname]
+    return enter, leave
+
+
+def dotrace(*args, **kwds):
+    """Function decorator to trace a function's entry and exit.
+
+    *args: categories in which to trace this function. Example usage:
+
+    @trace
+    def function(...):...
+
+    @trace('mycategory')
+    def function(...):...
+
+
+    """
+
+    recursive = kwds.get("recursive", False)
+
+    def decorator(func):
+        spec = None
+        logger = logging.getLogger("trace")
+
+        def wrapper(*args, **kwds):
+            if not logger.isEnabledFor(logging.INFO) or tls.tracing:
+                return func(*args, **kwds)
+
+            fname, ftype = find_function_info(func, spec, args)
+
+            try:
+                tls.tracing = True
+                enter, leave = create_events(fname, spec, args, kwds)
+
+                try:
+                    logger.info("".join(enter))
+                    tls.indent += 1
+                    try:
+                        try:
+                            tls.tracing = False
+                            result = func(*args, **kwds)
+                        finally:
+                            tls.tracing = True
+                    except:  # noqa: E722
+                        type, value, traceback = sys.exc_info()
+                        leave.append(" => exception thrown\n\traise ")
+                        mname = type.__module__
+                        if mname != "__main__":
+                            leave.append(mname)
+                            leave.append(".")
+                        leave.append(type.__name__)
+                        if value.args:
+                            leave.append("(")
+                            leave.append(", ".join(chop(v) for v in value.args))
+                            leave.append(")")
+                        else:
+                            leave.append("()")
+                        raise
+                    else:
+                        if result is not None:
+                            leave.append(" -> ")
+                            leave.append(chop(result))
+                finally:
+                    tls.indent -= 1
+                    logger.info("".join(leave))
+            finally:
+                tls.tracing = False
+            return result
+
+        # wrapper end
+
+        rewrap = lambda x: x
+        # Unwrap already wrapped functions
+        # (to be rewrapped again later)
+        if isinstance(func, classmethod):
+            rewrap = type(func)
+            # Note: 'func.__func__' only works in Python 3
+            func = func.__get__(True).__func__
+        elif isinstance(func, staticmethod):
+            rewrap = type(func)
+            # Note: 'func.__func__' only works in Python 3
+            func = func.__get__(True)
+        elif isinstance(func, property):
+            raise NotImplementedError
+
+        spec = inspect.getfullargspec(func)
+        return rewrap(wraps(func)(wrapper))
+
+    arg0 = len(args) and args[0] or None
+    # not supported yet...
+    if recursive:
+        raise NotImplementedError
+        if inspect.ismodule(arg0):
+            for n, f in inspect.getmembers(arg0, inspect.isfunction):
+                setattr(arg0, n, decorator(f))
+            for n, c in inspect.getmembers(arg0, inspect.isclass):
+                dotrace(c, *args, recursive=recursive)
+        elif inspect.isclass(arg0):
+            for n, f in inspect.getmembers(
+                arg0, lambda x: (inspect.isfunction(x) or inspect.ismethod(x))
+            ):
+                setattr(arg0, n, decorator(f))
+
+    if callable(arg0) or type(arg0) in (classmethod, staticmethod):
+        return decorator(arg0)
+    elif isinstance(arg0, property):
+        # properties combine up to three functions: 'get', 'set', 'del',
+        # so let's wrap them all.
+        pget, pset, pdel = None, None, None
+        if arg0.fget:
+            pget = decorator(arg0.fget)
+        if arg0.fset:
+            pset = decorator(arg0.fset)
+        if arg0.fdel:
+            pdel = decorator(arg0.fdel)
+        return property(pget, pset, pdel)
+
+    else:
+        return decorator
+
+
+def notrace(*args, **kwds):
+    """Just a no-op in case tracing is disabled."""
+
+    def decorator(func):
+        return func
+
+    arg0 = len(args) and args[0] or None
+
+    if callable(arg0) or type(arg0) in (classmethod, staticmethod):
+        return decorator(arg0)
+    else:
+        return decorator
+
+
+def doevent(msg):
+    msg = ["== ", tls.indent * " ", msg]
+    logger = logging.getLogger("trace")
+    logger.info("".join(msg))
+
+
+def noevent(msg):
+    pass
+
+
+if config.TRACE:
+    logger = logging.getLogger("trace")
+    logger.setLevel(logging.INFO)
+    logger.handlers = [logging.StreamHandler()]
+    trace = dotrace
+    event = doevent
+else:
+    trace = notrace
+    event = noevent