numba-cuda 0.22.1__cp311-cp311-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl

numba_cuda/numba/cuda/core/optional.py

@@ -0,0 +1,129 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import operator
+
+from numba.cuda import types, typing
+from numba.cuda import cgutils
+
+from numba.cuda.core.imputils import Registry, impl_ret_untracked
+
+registry = Registry("optional")
+lower_builtin = registry.lower
+lower_cast = registry.lower_cast
+lower_getattr_generic = registry.lower_getattr_generic
+lower_setattr_generic = registry.lower_setattr_generic
+
+
+def always_return_true_impl(context, builder, sig, args):
+    return cgutils.true_bit
+
+
+def always_return_false_impl(context, builder, sig, args):
+    return cgutils.false_bit
+
+
+def optional_is_none(context, builder, sig, args):
+    """
+    Check if an Optional value is invalid
+    """
+    [lty, rty] = sig.args
+    [lval, rval] = args
+
+    # Make sure None is on the right
+    if lty == types.none:
+        lty, rty = rty, lty
+        lval, rval = rval, lval
+
+    opt_type = lty
+    opt_val = lval
+
+    opt = context.make_helper(builder, opt_type, opt_val)
+    res = builder.not_(cgutils.as_bool_bit(builder, opt.valid))
+    return impl_ret_untracked(context, builder, sig.return_type, res)
+
+
+# None is/not None
+lower_builtin(operator.is_, types.none, types.none)(always_return_true_impl)
+
+# Optional is None
+lower_builtin(operator.is_, types.Optional, types.none)(optional_is_none)
+lower_builtin(operator.is_, types.none, types.Optional)(optional_is_none)
+
+
+@lower_getattr_generic(types.Optional)
+def optional_getattr(context, builder, typ, value, attr):
+    """
+    Optional.__getattr__ => redirect to the wrapped type.
+    """
+    inner_type = typ.type
+    val = context.cast(builder, value, typ, inner_type)
+    imp = context.get_getattr(inner_type, attr)
+    return imp(context, builder, inner_type, val, attr)
+
+
+@lower_setattr_generic(types.Optional)
+def optional_setattr(context, builder, sig, args, attr):
+    """
+    Optional.__setattr__ => redirect to the wrapped type.
+    """
+    basety, valty = sig.args
+    target, val = args
+    target_type = basety.type
+    target = context.cast(builder, target, basety, target_type)
+
+    newsig = typing.signature(sig.return_type, target_type, valty)
+    imp = context.get_setattr(attr, newsig)
+    return imp(builder, (target, val))
+
+
+@lower_cast(types.Optional, types.Optional)
+def optional_to_optional(context, builder, fromty, toty, val):
+    """
+    The handling of optional->optional cast must be special cased for
+    correct propagation of None value.  Given type T and U. casting of
+    T? to U? (? denotes optional) should always succeed.   If the from-value
+    is None, the None value the casted value (U?) should be None; otherwise,
+    the from-value is casted to U. This is different from casting T? to U,
+    which requires the from-value must not be None.
+    """
+    optval = context.make_helper(builder, fromty, value=val)
+    validbit = cgutils.as_bool_bit(builder, optval.valid)
+    # Create uninitialized optional value
+    outoptval = context.make_helper(builder, toty)
+
+    with builder.if_else(validbit) as (is_valid, is_not_valid):
+        with is_valid:
+            # Cast internal value
+            outoptval.valid = cgutils.true_bit
+            outoptval.data = context.cast(
+                builder, optval.data, fromty.type, toty.type
+            )
+
+        with is_not_valid:
+            # Store None to result
+            outoptval.valid = cgutils.false_bit
+            outoptval.data = cgutils.get_null_value(outoptval.data.type)
+
+    return outoptval._getvalue()
+
+
+@lower_cast(types.Any, types.Optional)
+def any_to_optional(context, builder, fromty, toty, val):
+    if fromty == types.none:
+        return context.make_optional_none(builder, toty.type)
+    else:
+        val = context.cast(builder, val, fromty, toty.type)
+        return context.make_optional_value(builder, toty.type, val)
+
+
+@lower_cast(types.Optional, types.Any)
+@lower_cast(types.Optional, types.Boolean)
+def optional_to_any(context, builder, fromty, toty, val):
+    optval = context.make_helper(builder, fromty, value=val)
+    validbit = cgutils.as_bool_bit(builder, optval.valid)
+    with builder.if_then(builder.not_(validbit), likely=False):
+        msg = "expected %s, got None" % (fromty.type,)
+        context.call_conv.return_user_exc(builder, TypeError, (msg,))
+
+    return context.cast(builder, optval.data, fromty.type, toty)

numba_cuda/numba/cuda/core/options.py

@@ -0,0 +1,262 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+"""
+Defines CUDA Options for use in the CUDA target
+"""
+
+from abc import ABCMeta, abstractmethod
+
+
+class AbstractOptionValue(metaclass=ABCMeta):
+    """Abstract base class for custom option values."""
+
+    @abstractmethod
+    def encode(self) -> str:
+        """Returns an encoding of the values"""
+        ...
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}({self.encode()})"
+
+
+class FastMathOptions(AbstractOptionValue):
+    """
+    Options for controlling fast math optimization.
+    """
+
+    def __init__(self, value):
+        # https://releases.llvm.org/7.0.0/docs/LangRef.html#fast-math-flags
+        valid_flags = {
+            "fast",
+            "nnan",
+            "ninf",
+            "nsz",
+            "arcp",
+            "contract",
+            "afn",
+            "reassoc",
+        }
+
+        if isinstance(value, FastMathOptions):
+            self.flags = value.flags.copy()
+        elif value is True:
+            self.flags = {"fast"}
+        elif value is False:
+            self.flags = set()
+        elif isinstance(value, set):
+            invalid = value - valid_flags
+            if invalid:
+                raise ValueError("Unrecognized fastmath flags: %s" % invalid)
+            self.flags = value
+        elif isinstance(value, dict):
+            invalid = set(value.keys()) - valid_flags
+            if invalid:
+                raise ValueError("Unrecognized fastmath flags: %s" % invalid)
+            self.flags = {v for v, enable in value.items() if enable}
+        else:
+            msg = "Expected fastmath option(s) to be either a bool, dict or set"
+            raise ValueError(msg)
+
+    def __bool__(self):
+        return bool(self.flags)
+
+    __nonzero__ = __bool__
+
+    def encode(self) -> str:
+        return str(self.flags)
+
+    def __eq__(self, other):
+        if type(other) is type(self):
+            return self.flags == other.flags
+        return NotImplemented
+
+
+class ParallelOptions(AbstractOptionValue):
+    """
+    Options for controlling auto parallelization.
+    """
+
+    __slots__ = (
+        "enabled",
+        "comprehension",
+        "reduction",
+        "inplace_binop",
+        "setitem",
+        "numpy",
+        "stencil",
+        "fusion",
+        "prange",
+    )
+
+    def __init__(self, value):
+        if isinstance(value, bool):
+            self.enabled = value
+            self.comprehension = value
+            self.reduction = value
+            self.inplace_binop = value
+            self.setitem = value
+            self.numpy = value
+            self.stencil = value
+            self.fusion = value
+            self.prange = value
+        elif isinstance(value, dict):
+            self.enabled = True
+            self.comprehension = value.pop("comprehension", True)
+            self.reduction = value.pop("reduction", True)
+            self.inplace_binop = value.pop("inplace_binop", True)
+            self.setitem = value.pop("setitem", True)
+            self.numpy = value.pop("numpy", True)
+            self.stencil = value.pop("stencil", True)
+            self.fusion = value.pop("fusion", True)
+            self.prange = value.pop("prange", True)
+            if value:
+                msg = "Unrecognized parallel options: %s" % value.keys()
+                raise NameError(msg)
+        elif isinstance(value, ParallelOptions):
+            self.enabled = value.enabled
+            self.comprehension = value.comprehension
+            self.reduction = value.reduction
+            self.inplace_binop = value.inplace_binop
+            self.setitem = value.setitem
+            self.numpy = value.numpy
+            self.stencil = value.stencil
+            self.fusion = value.fusion
+            self.prange = value.prange
+        else:
+            msg = "Expect parallel option to be either a bool or a dict"
+            raise ValueError(msg)
+
+    def _get_values(self):
+        """Get values as dictionary."""
+        return {k: getattr(self, k) for k in self.__slots__}
+
+    def __eq__(self, other):
+        if type(other) is type(self):
+            return self._get_values() == other._get_values()
+        return NotImplemented
+
+    def encode(self) -> str:
+        return ", ".join(f"{k}={v}" for k, v in self._get_values().items())
+
+
+class InlineOptions(AbstractOptionValue):
+    """
+    Options for controlling inlining
+    """
+
+    def __init__(self, value):
+        ok = False
+        if isinstance(value, str):
+            if value in ("always", "never"):
+                ok = True
+        else:
+            ok = hasattr(value, "__call__")
+
+        if ok:
+            self._inline = value
+        else:
+            msg = (
+                "kwarg 'inline' must be one of the strings 'always' or "
+                "'never', or it can be a callable that returns True/False. "
+                "Found value %s" % value
+            )
+            raise ValueError(msg)
+
+    @property
+    def is_never_inline(self):
+        """
+        True if never inline
+        """
+        return self._inline == "never"
+
+    @property
+    def is_always_inline(self):
+        """
+        True if always inline
+        """
+        return self._inline == "always"
+
+    @property
+    def has_cost_model(self):
+        """
+        True if a cost model is provided
+        """
+        return not (self.is_always_inline or self.is_never_inline)
+
+    @property
+    def value(self):
+        """
+        The raw value
+        """
+        return self._inline
+
+    def __eq__(self, other):
+        if type(other) is type(self):
+            return self.value == other.value
+        return NotImplemented
+
+    def encode(self) -> str:
+        return repr(self._inline)
+
+
+class TargetOptions:
+    """Target options maps user options from decorators to the
+    ``numba.cuda.core.compiler.Flags`` used by lowering and target context.
+    """
+
+    class Mapping:
+        def __init__(self, flag_name, apply=lambda x: x):
+            self.flag_name = flag_name
+            self.apply = apply
+
+    def finalize(self, flags, options):
+        """Subclasses can override this method to make target specific
+        customizations of default flags.
+
+        Parameters
+        ----------
+        flags : Flags
+        options : dict
+        """
+        pass
+
+    @classmethod
+    def parse_as_flags(cls, flags, options):
+        """Parse target options defined in ``options`` and set ``flags``
+        accordingly.
+
+        Parameters
+        ----------
+        flags : Flags
+        options : dict
+        """
+        opt = cls()
+        opt._apply(flags, options)
+        opt.finalize(flags, options)
+        return flags
+
+    def _apply(self, flags, options):
+        # Find all Mapping instances in the class
+        mappings = {}
+        cls = type(self)
+        for k in dir(cls):
+            v = getattr(cls, k)
+            if isinstance(v, cls.Mapping):
+                mappings[k] = v
+
+        used = set()
+        for k, mapping in mappings.items():
+            if k in options:
+                v = mapping.apply(options[k])
+                setattr(flags, mapping.flag_name, v)
+                used.add(k)
+
+        unused = set(options) - used
+        if unused:
+            # Unread options?
+            m = (
+                f"Unrecognized options: {unused}. "
+                f"Known options are {mappings.keys()}"
+            )
+            raise KeyError(m)

numba_cuda/numba/cuda/core/postproc.py

@@ -0,0 +1,249 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+from functools import cached_property
+from numba.cuda.core import analysis, ir
+from numba.cuda.core import ir_utils, transforms
+
+
+class YieldPoint(object):
+    def __init__(self, block, inst):
+        assert isinstance(block, ir.Block)
+        assert isinstance(inst, ir.Yield)
+        self.block = block
+        self.inst = inst
+        self.live_vars = None
+        self.weak_live_vars = None
+
+
+class GeneratorInfo(object):
+    def __init__(self):
+        # { index: YieldPoint }
+        self.yield_points = {}
+        # Ordered list of variable names
+        self.state_vars = []
+
+    def get_yield_points(self):
+        """
+        Return an iterable of YieldPoint instances.
+        """
+        return self.yield_points.values()
+
+
+class VariableLifetime(object):
+    """
+    For lazily building information of variable lifetime
+    """
+
+    def __init__(self, blocks):
+        self._blocks = blocks
+
+    @cached_property
+    def cfg(self):
+        return analysis.compute_cfg_from_blocks(self._blocks)
+
+    @cached_property
+    def usedefs(self):
+        return analysis.compute_use_defs(self._blocks)
+
+    @cached_property
+    def livemap(self):
+        return analysis.compute_live_map(
+            self.cfg, self._blocks, self.usedefs.usemap, self.usedefs.defmap
+        )
+
+    @cached_property
+    def deadmaps(self):
+        return analysis.compute_dead_maps(
+            self.cfg, self._blocks, self.livemap, self.usedefs.defmap
+        )
+
+
+# other packages that define new nodes add calls for inserting dels
+# format: {type:function}
+ir_extension_insert_dels = {}
+
+
+class PostProcessor(object):
+    """
+    A post-processor for Numba IR.
+    """
+
+    def __init__(self, func_ir):
+        self.func_ir = func_ir
+
+    def run(self, emit_dels: bool = False, extend_lifetimes: bool = False):
+        """
+        Run the following passes over Numba IR:
+        - canonicalize the CFG
+        - emit explicit `del` instructions for variables
+        - compute lifetime of variables
+        - compute generator info (if function is a generator function)
+        """
+        self.func_ir.blocks = transforms.canonicalize_cfg(self.func_ir.blocks)
+        vlt = VariableLifetime(self.func_ir.blocks)
+        self.func_ir.variable_lifetime = vlt
+
+        bev = analysis.compute_live_variables(
+            vlt.cfg,
+            self.func_ir.blocks,
+            vlt.usedefs.defmap,
+            vlt.deadmaps.combined,
+        )
+        for offset, ir_block in self.func_ir.blocks.items():
+            self.func_ir.block_entry_vars[ir_block] = bev[offset]
+
+        if self.func_ir.is_generator:
+            self.func_ir.generator_info = GeneratorInfo()
+            self._compute_generator_info()
+        else:
+            self.func_ir.generator_info = None
+
+        # Emit del nodes, do this last as the generator info parsing generates
+        # and then strips dels as part of its analysis.
+        if emit_dels:
+            self._insert_var_dels(extend_lifetimes=extend_lifetimes)
+
+    def _populate_generator_info(self):
+        """
+        Fill `index` for the Yield instruction and create YieldPoints.
+        """
+        dct = self.func_ir.generator_info.yield_points
+        assert not dct, "rerunning _populate_generator_info"
+        for block in self.func_ir.blocks.values():
+            for inst in block.body:
+                if isinstance(inst, ir.Assign):
+                    yieldinst = inst.value
+                    if isinstance(yieldinst, ir.Yield):
+                        index = len(dct) + 1
+                        yieldinst.index = index
+                        yp = YieldPoint(block, yieldinst)
+                        dct[yieldinst.index] = yp
+
+    def _compute_generator_info(self):
+        """
+        Compute the generator's state variables as the union of live variables
+        at all yield points.
+        """
+        # generate del info, it's used in analysis here, strip it out at the end
+        self._insert_var_dels()
+        self._populate_generator_info()
+        gi = self.func_ir.generator_info
+        for yp in gi.get_yield_points():
+            live_vars = set(self.func_ir.get_block_entry_vars(yp.block))
+            weak_live_vars = set()
+            stmts = iter(yp.block.body)
+            for stmt in stmts:
+                if isinstance(stmt, ir.Assign):
+                    if stmt.value is yp.inst:
+                        break
+                    live_vars.add(stmt.target.name)
+                elif isinstance(stmt, ir.Del):
+                    live_vars.remove(stmt.value)
+            else:
+                assert 0, "couldn't find yield point"
+            # Try to optimize out any live vars that are deleted immediately
+            # after the yield point.
+            for stmt in stmts:
+                if isinstance(stmt, ir.Del):
+                    name = stmt.value
+                    if name in live_vars:
+                        live_vars.remove(name)
+                        weak_live_vars.add(name)
+                else:
+                    break
+            yp.live_vars = live_vars
+            yp.weak_live_vars = weak_live_vars
+
+        st = set()
+        for yp in gi.get_yield_points():
+            st |= yp.live_vars
+            st |= yp.weak_live_vars
+        gi.state_vars = sorted(st)
+        self.remove_dels()
+
+    def _insert_var_dels(self, extend_lifetimes=False):
+        """
+        Insert del statements for each variable.
+        Returns a 2-tuple of (variable definition map, variable deletion map)
+        which indicates variables defined and deleted in each block.
+
+        The algorithm avoids relying on explicit knowledge on loops and
+        distinguish between variables that are defined locally vs variables that
+        come from incoming blocks.
+        We start with simple usage (variable reference) and definition (variable
+        creation) maps on each block. Propagate the liveness info to predecessor
+        blocks until it stabilize, at which point we know which variables must
+        exist before entering each block. Then, we compute the end of variable
+        lives and insert del statements accordingly. Variables are deleted after
+        the last use. Variable referenced by terminators (e.g. conditional
+        branch and return) are deleted by the successors or the caller.
+        """
+        vlt = self.func_ir.variable_lifetime
+        self._patch_var_dels(
+            vlt.deadmaps.internal,
+            vlt.deadmaps.escaping,
+            extend_lifetimes=extend_lifetimes,
+        )
+
+    def _patch_var_dels(
+        self, internal_dead_map, escaping_dead_map, extend_lifetimes=False
+    ):
+        """
+        Insert delete in each block
+        """
+        for offset, ir_block in self.func_ir.blocks.items():
+            # for each internal var, insert delete after the last use
+            internal_dead_set = internal_dead_map[offset].copy()
+            delete_pts = []
+            # for each statement in reverse order
+            for stmt in reversed(ir_block.body[:-1]):
+                # internal vars that are used here
+                live_set = set(v.name for v in stmt.list_vars())
+                dead_set = live_set & internal_dead_set
+                for T, def_func in ir_extension_insert_dels.items():
+                    if isinstance(stmt, T):
+                        done_dels = def_func(stmt, dead_set)
+                        dead_set -= done_dels
+                        internal_dead_set -= done_dels
+                # used here but not afterwards
+                delete_pts.append((stmt, dead_set))
+                internal_dead_set -= dead_set
+
+            # rewrite body and insert dels
+            body = []
+            lastloc = ir_block.loc
+            del_store = []
+            for stmt, delete_set in reversed(delete_pts):
+                # If using extended lifetimes then the Dels are all put at the
+                # block end just ahead of the terminator, so associate their
+                # location with the terminator.
+                if extend_lifetimes:
+                    lastloc = ir_block.body[-1].loc
+                else:
+                    lastloc = stmt.loc
+                # Ignore dels (assuming no user inserted deletes)
+                if not isinstance(stmt, ir.Del):
+                    body.append(stmt)
+                # note: the reverse sort is not necessary for correctness
+                #       it is just to minimize changes to test for now
+                for var_name in sorted(delete_set, reverse=True):
+                    delnode = ir.Del(var_name, loc=lastloc)
+                    if extend_lifetimes:
+                        del_store.append(delnode)
+                    else:
+                        body.append(delnode)
+            if extend_lifetimes:
+                body.extend(del_store)
+            body.append(ir_block.body[-1])  # terminator
+            ir_block.body = body
+
+            # vars to delete at the start
+            escape_dead_set = escaping_dead_map[offset]
+            for var_name in sorted(escape_dead_set):
+                ir_block.prepend(ir.Del(var_name, loc=ir_block.body[0].loc))
+
+    def remove_dels(self):
+        """
+        Strips the IR of Del nodes
+        """
+        ir_utils.remove_dels(self.func_ir.blocks)