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

numba_cuda/numba/cuda/np/ufunc/decorators.py

@@ -0,0 +1,203 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+from numba.cuda.core.registry import DelayedRegistry
+from numba.cuda.vectorizers import CUDAVectorize, CUDAGUFuncVectorize
+
+
+class _BaseVectorize(object):
+    @classmethod
+    def get_identity(cls, kwargs):
+        return kwargs.pop("identity", None)
+
+    @classmethod
+    def get_cache(cls, kwargs):
+        return kwargs.pop("cache", False)
+
+    @classmethod
+    def get_writable_args(cls, kwargs):
+        return kwargs.pop("writable_args", ())
+
+    @classmethod
+    def get_target_implementation(cls, kwargs):
+        target = kwargs.pop("target", "cpu")
+        try:
+            return cls.target_registry[target]
+        except KeyError:
+            raise ValueError("Unsupported target: %s" % target)
+
+
+class Vectorize(_BaseVectorize):
+    target_registry = DelayedRegistry({"cuda": CUDAVectorize})
+
+    def __new__(cls, func, **kws):
+        identity = cls.get_identity(kws)
+        cache = cls.get_cache(kws)
+        imp = cls.get_target_implementation(kws)
+        return imp(func, identity=identity, cache=cache, targetoptions=kws)
+
+
+class GUVectorize(_BaseVectorize):
+    target_registry = DelayedRegistry({"cuda": CUDAGUFuncVectorize})
+
+    def __new__(cls, func, signature, **kws):
+        identity = cls.get_identity(kws)
+        cache = cls.get_cache(kws)
+        imp = cls.get_target_implementation(kws)
+        writable_args = cls.get_writable_args(kws)
+        return imp(
+            func,
+            signature,
+            identity=identity,
+            cache=cache,
+            targetoptions=kws,
+            writable_args=writable_args,
+        )
+
+
+def vectorize(ftylist_or_function=(), target="cuda", **kws):
+    """vectorize(ftylist_or_function=(), target='cuda', identity=None, **kws)
+
+    A decorator that creates a NumPy ufunc object using Numba compiled
+    code.  When no arguments or only keyword arguments are given,
+    vectorize will return a Numba dynamic ufunc (DUFunc) object, where
+    compilation/specialization may occur at call-time.
+
+    Args
+    -----
+    ftylist_or_function: function or iterable
+
+        When the first argument is a function, signatures are dealt
+        with at call-time.
+
+        When the first argument is an iterable of type signatures,
+        which are either function type object or a string describing
+        the function type, signatures are finalized at decoration
+        time.
+
+    Keyword Args
+    ------------
+
+    target: str
+            A string for code generation target.  Default to "cuda".
+
+    identity: int, str, or None
+        The identity (or unit) value for the element-wise function
+        being implemented.  Allowed values are None (the default), 0, 1,
+        and "reorderable".
+
+    cache: bool
+        Turns on caching.
+
+
+    Returns
+    --------
+
+    A NumPy universal function
+
+    Examples
+    -------
+        @vectorize(['float32(float32, float32)',
+                    'float64(float64, float64)'], identity=0)
+        def sum(a, b):
+            return a + b
+
+        @vectorize
+        def sum(a, b):
+            return a + b
+
+        @vectorize(identity=1)
+        def mul(a, b):
+            return a * b
+
+    """
+    if isinstance(ftylist_or_function, str):
+        # Common user mistake
+        ftylist = [ftylist_or_function]
+    elif ftylist_or_function is not None:
+        ftylist = ftylist_or_function
+
+    def wrap(func):
+        kws["target"] = target
+        vec = Vectorize(func, **kws)
+        for sig in ftylist:
+            vec.add(sig)
+        if len(ftylist) > 0:
+            vec.disable_compile()
+        return vec.build_ufunc()
+
+    return wrap
+
+
+def guvectorize(*args, **kwargs):
+    """guvectorize(ftylist, signature, target='cuda', identity=None, **kws)
+
+    A decorator to create NumPy generalized-ufunc object from Numba compiled
+    code.
+
+    Args
+    -----
+    ftylist: iterable
+        An iterable of type signatures, which are either
+        function type object or a string describing the
+        function type.
+
+    signature: str
+        A NumPy generalized-ufunc signature.
+        e.g. "(m, n), (n, p)->(m, p)"
+
+    identity: int, str, or None
+        The identity (or unit) value for the element-wise function
+        being implemented.  Allowed values are None (the default), 0, 1,
+        and "reorderable".
+
+    cache: bool
+        Turns on caching.
+
+    writable_args: tuple
+        a tuple of indices of input variables that are writable.
+
+    target: str
+            A string for code generation target.  Defaults to "cuda".
+
+    Returns
+    --------
+
+    A NumPy generalized universal-function
+
+    Example
+    -------
+        @guvectorize(['void(int32[:,:], int32[:,:], int32[:,:])',
+                      'void(float32[:,:], float32[:,:], float32[:,:])'],
+                      '(x, y),(x, y)->(x, y)')
+        def add_2d_array(a, b, c):
+            for i in range(c.shape[0]):
+                for j in range(c.shape[1]):
+                    c[i, j] = a[i, j] + b[i, j]
+
+    """
+    if len(args) == 1:
+        ftylist = []
+        signature = args[0]
+        kwargs.setdefault("is_dynamic", True)
+    elif len(args) == 2:
+        ftylist = args[0]
+        signature = args[1]
+    else:
+        raise TypeError("guvectorize() takes one or two positional arguments")
+
+    if isinstance(ftylist, str):
+        # Common user mistake
+        ftylist = [ftylist]
+
+    kwargs.setdefault("target", "cuda")
+
+    def wrap(func):
+        guvec = GUVectorize(func, signature, **kwargs)
+        for fty in ftylist:
+            guvec.add(fty)
+        if len(ftylist) > 0:
+            guvec.disable_compile()
+        return guvec.build_ufunc()
+
+    return wrap

numba_cuda/numba/cuda/np/ufunc/sigparse.py

@@ -0,0 +1,68 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import tokenize
+import string
+
+
+def parse_signature(sig):
+    """Parse generalized ufunc signature.
+
+    NOTE: ',' (COMMA) is a delimiter; not separator.
+          This means trailing comma is legal.
+    """
+
+    def stripws(s):
+        return "".join(c for c in s if c not in string.whitespace)
+
+    def tokenizer(src):
+        def readline():
+            yield src
+
+        gen = readline()
+        return tokenize.generate_tokens(lambda: next(gen))
+
+    def parse(src):
+        tokgen = tokenizer(src)
+        while True:
+            tok = next(tokgen)
+            if tok[1] == "(":
+                symbols = []
+                while True:
+                    tok = next(tokgen)
+                    if tok[1] == ")":
+                        break
+                    elif tok[0] == tokenize.NAME:
+                        symbols.append(tok[1])
+                    elif tok[1] == ",":
+                        continue
+                    else:
+                        raise ValueError('bad token in signature "%s"' % tok[1])
+                yield tuple(symbols)
+                tok = next(tokgen)
+                if tok[1] == ",":
+                    continue
+                elif tokenize.ISEOF(tok[0]):
+                    break
+            elif tokenize.ISEOF(tok[0]):
+                break
+            else:
+                raise ValueError('bad token in signature "%s"' % tok[1])
+
+    ins, _, outs = stripws(sig).partition("->")
+    inputs = list(parse(ins))
+    outputs = list(parse(outs))
+
+    # check that all output symbols are defined in the inputs
+    isym = set()
+    osym = set()
+    for grp in inputs:
+        isym |= set(grp)
+    for grp in outputs:
+        osym |= set(grp)
+
+    diff = osym.difference(isym)
+    if diff:
+        raise NameError("undefined output symbols: %s" % ",".join(sorted(diff)))
+
+    return inputs, outputs

numba_cuda/numba/cuda/np/ufunc/ufuncbuilder.py

@@ -0,0 +1,65 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+from numba.cuda.core import sigutils
+
+# Utility functions
+
+# HACK: These are explicitly defined here to avoid having a CExt just to import these constants.
+#   np doesn't expose these in the python API.
+PyUFunc_Zero = 0
+PyUFunc_One = 1
+PyUFunc_None = -1
+PyUFunc_ReorderableNone = -2
+
+
+def _compile_element_wise_function(nb_func, targetoptions, sig):
+    # Do compilation
+    # Return CompileResult to test
+    cres = nb_func.compile(sig, **targetoptions)
+    args, return_type = sigutils.normalize_signature(sig)
+    return cres, args, return_type
+
+
+# Class definitions
+
+
+class _BaseUFuncBuilder(object):
+    def add(self, sig=None):
+        if hasattr(self, "targetoptions"):
+            targetoptions = self.targetoptions
+        else:
+            targetoptions = self.nb_func.targetoptions
+        cres, args, return_type = _compile_element_wise_function(
+            self.nb_func, targetoptions, sig
+        )
+        sig = self._finalize_signature(cres, args, return_type)
+        self._sigs.append(sig)
+        self._cres[sig] = cres
+        return cres
+
+    def disable_compile(self):
+        """
+        Disable the compilation of new signatures at call time.
+        """
+        # Override this for implementations that support lazy compilation
+
+
+_identities = {
+    0: PyUFunc_Zero,
+    1: PyUFunc_One,
+    None: PyUFunc_None,
+    "reorderable": PyUFunc_ReorderableNone,
+}
+
+
+def parse_identity(identity):
+    """
+    Parse an identity value and return the corresponding low-level value
+    for Numpy.
+    """
+    try:
+        identity = _identities[identity]
+    except KeyError:
+        raise ValueError("Invalid identity value %r" % (identity,))
+    return identity