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

numba_cuda/numba/cuda/misc/llvm_pass_timings.py

@@ -0,0 +1,412 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import re
+import operator
+import heapq
+from collections import namedtuple
+from collections.abc import Sequence
+from contextlib import contextmanager
+from functools import cached_property
+
+from numba.cuda import config
+
+import llvmlite.binding as llvm
+
+
+class RecordLLVMPassTimings:
+    """A helper context manager to track LLVM pass timings."""
+
+    __slots__ = ["_data"]
+
+    def __enter__(self):
+        """Enables the pass timing in LLVM."""
+        llvm.set_time_passes(True)
+        return self
+
+    def __exit__(self, exc_val, exc_type, exc_tb):
+        """Reset timings and save report internally."""
+        self._data = llvm.report_and_reset_timings()
+        llvm.set_time_passes(False)
+        return
+
+    def get(self):
+        """Retrieve timing data for processing.
+
+        Returns
+        -------
+        timings: ProcessedPassTimings
+        """
+        return ProcessedPassTimings(self._data)
+
+
+PassTimingRecord = namedtuple(
+    "PassTimingRecord",
+    [
+        "user_time",
+        "user_percent",
+        "system_time",
+        "system_percent",
+        "user_system_time",
+        "user_system_percent",
+        "wall_time",
+        "wall_percent",
+        "pass_name",
+        "instruction",
+    ],
+)
+
+
+def _adjust_timings(records):
+    """Adjust timing records because of truncated information.
+
+    Details: The percent information can be used to improve the timing
+    information.
+
+    Returns
+    -------
+    res: List[PassTimingRecord]
+    """
+    total_rec = records[-1]
+    assert total_rec.pass_name == "Total"  # guard for implementation error
+
+    def make_adjuster(attr):
+        time_attr = f"{attr}_time"
+        percent_attr = f"{attr}_percent"
+        time_getter = operator.attrgetter(time_attr)
+
+        def adjust(d):
+            """Compute percent x total_time = adjusted"""
+            total = time_getter(total_rec)
+            adjusted = total * d[percent_attr] * 0.01
+            d[time_attr] = adjusted
+            return d
+
+        return adjust
+
+    # Make adjustment functions for each field
+    adj_fns = [
+        make_adjuster(x) for x in ["user", "system", "user_system", "wall"]
+    ]
+
+    # Extract dictionaries from the namedtuples
+    dicts = map(lambda x: x._asdict(), records)
+
+    def chained(d):
+        # Chain the adjustment functions
+        for fn in adj_fns:
+            d = fn(d)
+        # Reconstruct the namedtuple
+        return PassTimingRecord(**d)
+
+    return list(map(chained, dicts))
+
+
+class ProcessedPassTimings:
+    """A class for processing raw timing report from LLVM.
+
+    The processing is done lazily so we don't waste time processing unused
+    timing information.
+    """
+
+    def __init__(self, raw_data):
+        self._raw_data = raw_data
+
+    def __bool__(self):
+        return bool(self._raw_data)
+
+    def get_raw_data(self):
+        """Returns the raw string data.
+
+        Returns
+        -------
+        res: str
+        """
+        return self._raw_data
+
+    def get_total_time(self):
+        """Compute the total time spend in all passes.
+
+        Returns
+        -------
+        res: float
+        """
+        return self.list_records()[-1].wall_time
+
+    def list_records(self):
+        """Get the processed data for the timing report.
+
+        Returns
+        -------
+        res: List[PassTimingRecord]
+        """
+        return self._processed
+
+    def list_top(self, n):
+        """Returns the top(n) most time-consuming (by wall-time) passes.
+
+        Parameters
+        ----------
+        n: int
+            This limits the maximum number of items to show.
+            This function will show the ``n`` most time-consuming passes.
+
+        Returns
+        -------
+        res: List[PassTimingRecord]
+            Returns the top(n) most time-consuming passes in descending order.
+        """
+        records = self.list_records()
+        key = operator.attrgetter("wall_time")
+        return heapq.nlargest(n, records[:-1], key)
+
+    def summary(self, topn=5, indent=0):
+        """Return a string summarizing the timing information.
+
+        Parameters
+        ----------
+        topn: int; optional
+            This limits the maximum number of items to show.
+            This function will show the ``topn`` most time-consuming passes.
+        indent: int; optional
+            Set the indentation level. Defaults to 0 for no indentation.
+
+        Returns
+        -------
+        res: str
+        """
+        buf = []
+        prefix = " " * indent
+
+        def ap(arg):
+            buf.append(f"{prefix}{arg}")
+
+        ap(f"Total {self.get_total_time():.4f}s")
+        ap("Top timings:")
+        for p in self.list_top(topn):
+            ap(f"  {p.wall_time:.4f}s ({p.wall_percent:5}%) {p.pass_name}")
+        return "\n".join(buf)
+
+    @cached_property
+    def _processed(self):
+        """A cached property for lazily processing the data and returning it.
+
+        See ``_process()`` for details.
+        """
+        return self._process()
+
+    def _process(self):
+        """Parses the raw string data from LLVM timing report and attempts
+        to improve the data by recomputing the times
+        (See `_adjust_timings()``).
+        """
+
+        def parse(raw_data):
+            """A generator that parses the raw_data line-by-line to extract
+            timing information for each pass.
+            """
+            lines = raw_data.splitlines()
+            colheader = r"[a-zA-Z+ ]+"
+            # Take at least one column header.
+            multicolheaders = rf"(?:\s*-+{colheader}-+)+"
+
+            line_iter = iter(lines)
+            # find column headers
+            header_map = {
+                "User Time": "user",
+                "System Time": "system",
+                "User+System": "user_system",
+                "Wall Time": "wall",
+                "Instr": "instruction",
+                "Name": "pass_name",
+            }
+            for ln in line_iter:
+                m = re.match(multicolheaders, ln)
+                if m:
+                    # Get all the column headers
+                    raw_headers = re.findall(r"[a-zA-Z][a-zA-Z+ ]+", ln)
+                    headers = [header_map[k.strip()] for k in raw_headers]
+                    break
+
+            assert headers[-1] == "pass_name"
+            # compute the list of available attributes from the column headers
+            attrs = []
+            n = r"\s*((?:[0-9]+\.)?[0-9]+)"
+            pat = ""
+            for k in headers[:-1]:
+                if k == "instruction":
+                    pat += n
+                else:
+                    attrs.append(f"{k}_time")
+                    attrs.append(f"{k}_percent")
+                    pat += rf"\s+(?:{n}\s*\({n}%\)|-+)"
+
+            # put default value 0.0 to all missing attributes
+            missing = {}
+            for k in PassTimingRecord._fields:
+                if k not in attrs and k != "pass_name":
+                    missing[k] = 0.0
+            # parse timings
+            pat += r"\s*(.*)"
+            for ln in line_iter:
+                m = re.match(pat, ln)
+                if m is not None:
+                    raw_data = list(m.groups())
+                    data = {
+                        k: float(v) if v is not None else 0.0
+                        for k, v in zip(attrs, raw_data)
+                    }
+                    data.update(missing)
+                    pass_name = raw_data[-1]
+                    rec = PassTimingRecord(
+                        pass_name=pass_name,
+                        **data,
+                    )
+                    yield rec
+                    if rec.pass_name == "Total":
+                        # "Total" means the report has ended
+                        break
+            # Check that we have reach the end of the report
+            remaining = "\n".join(line_iter)
+            if remaining:
+                raise ValueError(
+                    f"unexpected text after parser finished:\n{remaining}"
+                )
+
+        # Parse raw data
+        records = list(parse(self._raw_data))
+        return _adjust_timings(records)
+
+
+NamedTimings = namedtuple("NamedTimings", ["name", "timings"])
+
+
+class PassTimingsCollection(Sequence):
+    """A collection of pass timings.
+
+    This class implements the ``Sequence`` protocol for accessing the
+    individual timing records.
+    """
+
+    def __init__(self, name):
+        self._name = name
+        self._records = []
+
+    @contextmanager
+    def record(self, name):
+        """Record new timings and append to this collection.
+
+        Note: this is mainly for internal use inside the compiler pipeline.
+
+        See also ``RecordLLVMPassTimings``
+
+        Parameters
+        ----------
+        name: str
+            Name for the records.
+        """
+        if config.LLVM_PASS_TIMINGS:
+            # Recording of pass timings is enabled
+            with RecordLLVMPassTimings() as timings:
+                yield
+            rec = timings.get()
+            # Only keep non-empty records
+            if rec:
+                self._append(name, rec)
+        else:
+            # Do nothing. Recording of pass timings is disabled.
+            yield
+
+    def _append(self, name, timings):
+        """Append timing records
+
+        Parameters
+        ----------
+        name: str
+            Name for the records.
+        timings: ProcessedPassTimings
+            the timing records.
+        """
+        self._records.append(NamedTimings(name, timings))
+
+    def get_total_time(self):
+        """Computes the sum of the total time across all contained timings.
+
+        Returns
+        -------
+        res: float or None
+            Returns the total number of seconds or None if no timings were
+            recorded
+        """
+        if self._records:
+            return sum(r.timings.get_total_time() for r in self._records)
+        else:
+            return None
+
+    def list_longest_first(self):
+        """Returns the timings in descending order of total time duration.
+
+        Returns
+        -------
+        res: List[ProcessedPassTimings]
+        """
+        return sorted(
+            self._records,
+            key=lambda x: x.timings.get_total_time(),
+            reverse=True,
+        )
+
+    @property
+    def is_empty(self):
+        """ """
+        return not self._records
+
+    def summary(self, topn=5):
+        """Return a string representing the summary of the timings.
+
+        Parameters
+        ----------
+        topn: int; optional, default=5.
+            This limits the maximum number of items to show.
+            This function will show the ``topn`` most time-consuming passes.
+
+        Returns
+        -------
+        res: str
+
+        See also ``ProcessedPassTimings.summary()``
+        """
+        if self.is_empty:
+            return "No pass timings were recorded"
+        else:
+            buf = []
+            ap = buf.append
+            ap(f"Printing pass timings for {self._name}")
+            overall_time = self.get_total_time()
+            ap(f"Total time: {overall_time:.4f}")
+            for i, r in enumerate(self._records):
+                ap(f"== #{i} {r.name}")
+                percent = r.timings.get_total_time() / overall_time * 100
+                ap(f" Percent: {percent:.1f}%")
+                ap(r.timings.summary(topn=topn, indent=1))
+            return "\n".join(buf)
+
+    def __getitem__(self, i):
+        """Get the i-th timing record.
+
+        Returns
+        -------
+        res: (name, timings)
+            A named tuple with two fields:
+
+            - name: str
+            - timings: ProcessedPassTimings
+        """
+        return self._records[i]
+
+    def __len__(self):
+        """Length of this collection."""
+        return len(self._records)
+
+    def __str__(self):
+        return self.summary()

numba_cuda/numba/cuda/misc/special.py

@@ -0,0 +1,94 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+
+class prange(object):
+    """Provides a 1D parallel iterator that generates a sequence of integers.
+    In non-parallel contexts, prange is identical to range.
+    """
+
+    def __new__(cls, *args):
+        return range(*args)
+
+
+def _gdb_python_call_gen(func_name, *args):
+    # generates a call to a function containing a compiled in gdb command,
+    # this is to make `numba.gdb*` work in the interpreter.
+    import numba
+
+    fn = getattr(numba, func_name)
+    argstr = ",".join(['"%s"' for _ in args]) % args
+    defn = """def _gdb_func_injection():\n\t%s(%s)\n
+    """ % (func_name, argstr)
+    l = {}
+    exec(defn, {func_name: fn}, l)
+    return numba.njit(l["_gdb_func_injection"])
+
+
+def gdb(*args):
+    """
+    Calling this function will invoke gdb and attach it to the current process
+    at the call site. Arguments are strings in the gdb command language syntax
+    which will be executed by gdb once initialisation has occurred.
+    """
+    _gdb_python_call_gen("gdb", *args)()
+
+
+def gdb_breakpoint():
+    """
+    Calling this function will inject a breakpoint at the call site that is
+    recognised by both `gdb` and `gdb_init`, this is to allow breaking at
+    multiple points. gdb will stop in the user defined code just after the frame
+    employed by the breakpoint returns.
+    """
+    _gdb_python_call_gen("gdb_breakpoint")()
+
+
+def gdb_init(*args):
+    """
+    Calling this function will invoke gdb and attach it to the current process
+    at the call site, then continue executing the process under gdb's control.
+    Arguments are strings in the gdb command language syntax which will be
+    executed by gdb once initialisation has occurred.
+    """
+    _gdb_python_call_gen("gdb_init", *args)()
+
+
+def literally(obj):
+    """Forces Numba to interpret *obj* as an Literal value.
+
+    *obj* must be either a literal or an argument of the caller function, where
+    the argument must be bound to a literal. The literal requirement
+    propagates up the call stack.
+
+    This function is intercepted by the compiler to alter the compilation
+    behavior to wrap the corresponding function parameters as ``Literal``.
+    It has **no effect** outside of nopython-mode (interpreter, and objectmode).
+
+    The current implementation detects literal arguments in two ways:
+
+    1. Scans for uses of ``literally`` via a compiler pass.
+    2. ``literally`` is overloaded to raise ``numba.errors.ForceLiteralArg``
+       to signal the dispatcher to treat the corresponding parameter
+       differently. This mode is to support indirect use (via a function call).
+
+    The execution semantic of this function is equivalent to an identity
+    function.
+
+    See :ghfile:`numba/tests/test_literal_dispatch.py` for examples.
+    """
+    return obj
+
+
+def literal_unroll(container):
+    return container
+
+
+__all__ = [
+    "prange",
+    "gdb",
+    "gdb_breakpoint",
+    "gdb_init",
+    "literally",
+    "literal_unroll",
+]

numba_cuda/numba/cuda/models.py

@@ -0,0 +1,56 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import functools
+
+from llvmlite import ir
+
+from numba.cuda.datamodel.registry import DataModelManager, register
+from numba.cuda.datamodel import PrimitiveModel
+from numba.cuda.datamodel.models import StructModel
+from numba.cuda.extending import core_models as models
+from numba.cuda import types
+from numba.cuda.types.ext_types import Dim3, GridGroup, CUDADispatcher, Bfloat16
+
+
+cuda_data_manager = DataModelManager()
+
+register_model = functools.partial(register, cuda_data_manager)
+
+
+@register_model(Dim3)
+class Dim3Model(StructModel):
+    def __init__(self, dmm, fe_type):
+        members = [("x", types.int32), ("y", types.int32), ("z", types.int32)]
+        super().__init__(dmm, fe_type, members)
+
+
+@register_model(GridGroup)
+class GridGroupModel(models.PrimitiveModel):
+    def __init__(self, dmm, fe_type):
+        be_type = ir.IntType(64)
+        super().__init__(dmm, fe_type, be_type)
+
+
+@register_model(types.Float)
+class FloatModel(models.PrimitiveModel):
+    def __init__(self, dmm, fe_type):
+        if fe_type == types.float16:
+            be_type = ir.IntType(16)
+        elif fe_type == types.float32:
+            be_type = ir.FloatType()
+        elif fe_type == types.float64:
+            be_type = ir.DoubleType()
+        else:
+            raise NotImplementedError(fe_type)
+        super(FloatModel, self).__init__(dmm, fe_type, be_type)
+
+
+register_model(CUDADispatcher)(models.OpaqueModel)
+
+
+@register_model(Bfloat16)
+class _model___nv_bfloat16(PrimitiveModel):
+    def __init__(self, dmm, fe_type):
+        be_type = ir.IntType(16)
+        super(_model___nv_bfloat16, self).__init__(dmm, fe_type, be_type)