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

numba_cuda/numba/cuda/np/polynomial/__init__.py

@@ -0,0 +1,4 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+from .polynomial_core import registry  # noqa: F401

numba_cuda/numba/cuda/np/polynomial/polynomial_core.py

@@ -0,0 +1,242 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+from numba.cuda.extending import (
+    core_models,
+    register_model,
+    type_callable,
+    unbox,
+    NativeValue,
+    make_attribute_wrapper,
+    box,
+)
+from numba.cuda import types
+from numba.cuda import cgutils
+import warnings
+from numba.cuda.core.errors import (
+    NumbaExperimentalFeatureWarning,
+    NumbaValueError,
+)
+from numpy.polynomial.polynomial import Polynomial
+from contextlib import ExitStack
+import numpy as np
+from llvmlite import ir
+from numba.cuda.core.imputils import Registry
+
+registry = Registry("np.polynomial_core")
+lower = registry.lower
+
+
+@register_model(types.PolynomialType)
+class PolynomialModel(core_models.StructModel):
+    def __init__(self, dmm, fe_type):
+        members = [
+            ("coef", fe_type.coef),
+            ("domain", fe_type.domain),
+            ("window", fe_type.window),
+            # Introduced in NumPy 1.24, maybe leave it out for now
+            # ('symbol', types.string)
+        ]
+        super(PolynomialModel, self).__init__(dmm, fe_type, members)
+
+
+@type_callable(Polynomial)
+def type_polynomial(context):
+    def typer(coef, domain=None, window=None):
+        default_domain = types.Array(types.int64, 1, "C")
+        double_domain = types.Array(types.double, 1, "C")
+        default_window = types.Array(types.int64, 1, "C")
+        double_window = types.Array(types.double, 1, "C")
+        double_coef = types.Array(types.double, 1, "C")
+
+        warnings.warn(
+            "Polynomial class is experimental",
+            category=NumbaExperimentalFeatureWarning,
+        )
+
+        if isinstance(coef, types.Array) and all(
+            [a is None for a in (domain, window)]
+        ):
+            if coef.ndim == 1:
+                # If Polynomial(coef) is called, coef is cast to double dtype,
+                # and domain and window are set to equal [-1, 1], i.e. have
+                # integer dtype
+                return types.PolynomialType(
+                    double_coef, default_domain, default_window, 1
+                )
+            else:
+                msg = "Coefficient array is not 1-d"
+                raise NumbaValueError(msg)
+        elif all([isinstance(a, types.Array) for a in (coef, domain, window)]):
+            if coef.ndim == 1:
+                if all([a.ndim == 1 for a in (domain, window)]):
+                    # If Polynomial(coef, domain, window) is called, then coef,
+                    # domain and window are cast to double dtype
+                    return types.PolynomialType(
+                        double_coef, double_domain, double_window, 3
+                    )
+            else:
+                msg = "Coefficient array is not 1-d"
+                raise NumbaValueError(msg)
+
+    return typer
+
+
+make_attribute_wrapper(types.PolynomialType, "coef", "coef")
+make_attribute_wrapper(types.PolynomialType, "domain", "domain")
+make_attribute_wrapper(types.PolynomialType, "window", "window")
+# Introduced in NumPy 1.24, maybe leave it out for now
+# make_attribute_wrapper(types.PolynomialType, 'symbol', 'symbol')
+
+
+@lower(Polynomial, types.Array)
+def impl_polynomial1(context, builder, sig, args):
+    def to_double(arr):
+        return np.asarray(arr, dtype=np.double)
+
+    def const_impl():
+        return np.asarray([-1, 1])
+
+    typ = sig.return_type
+    polynomial = cgutils.create_struct_proxy(typ)(context, builder)
+    sig_coef = sig.args[0].copy(dtype=types.double)(sig.args[0])
+    coef_cast = context.compile_internal(builder, to_double, sig_coef, args)
+    sig_domain = sig.args[0].copy(dtype=types.intp)()
+    sig_window = sig.args[0].copy(dtype=types.intp)()
+    domain_cast = context.compile_internal(builder, const_impl, sig_domain, ())
+    window_cast = context.compile_internal(builder, const_impl, sig_window, ())
+    polynomial.coef = coef_cast
+    polynomial.domain = domain_cast
+    polynomial.window = window_cast
+
+    return polynomial._getvalue()
+
+
+@lower(Polynomial, types.Array, types.Array, types.Array)
+def impl_polynomial3(context, builder, sig, args):
+    def to_double(coef):
+        return np.asarray(coef, dtype=np.double)
+
+    typ = sig.return_type
+    polynomial = cgutils.create_struct_proxy(typ)(context, builder)
+
+    coef_sig = sig.args[0].copy(dtype=types.double)(sig.args[0])
+    domain_sig = sig.args[1].copy(dtype=types.double)(sig.args[1])
+    window_sig = sig.args[2].copy(dtype=types.double)(sig.args[2])
+    coef_cast = context.compile_internal(
+        builder, to_double, coef_sig, (args[0],)
+    )
+    domain_cast = context.compile_internal(
+        builder, to_double, domain_sig, (args[1],)
+    )
+    window_cast = context.compile_internal(
+        builder, to_double, window_sig, (args[2],)
+    )
+
+    domain_helper = context.make_helper(
+        builder, domain_sig.return_type, value=domain_cast
+    )
+    window_helper = context.make_helper(
+        builder, window_sig.return_type, value=window_cast
+    )
+
+    i64 = ir.IntType(64)
+    two = i64(2)
+
+    s1 = builder.extract_value(domain_helper.shape, 0)
+    s2 = builder.extract_value(window_helper.shape, 0)
+    pred1 = builder.icmp_signed("!=", s1, two)
+    pred2 = builder.icmp_signed("!=", s2, two)
+
+    with cgutils.if_unlikely(builder, pred1):
+        context.call_conv.return_user_exc(
+            builder, ValueError, ("Domain has wrong number of elements.",)
+        )
+
+    with cgutils.if_unlikely(builder, pred2):
+        context.call_conv.return_user_exc(
+            builder, ValueError, ("Window has wrong number of elements.",)
+        )
+
+    polynomial.coef = coef_cast
+    polynomial.domain = domain_helper._getvalue()
+    polynomial.window = window_helper._getvalue()
+
+    return polynomial._getvalue()
+
+
+@unbox(types.PolynomialType)
+def unbox_polynomial(typ, obj, c):
+    """
+    Convert a Polynomial object to a native polynomial structure.
+    """
+    is_error_ptr = cgutils.alloca_once_value(c.builder, cgutils.false_bit)
+    polynomial = cgutils.create_struct_proxy(typ)(c.context, c.builder)
+    with ExitStack() as stack:
+        natives = []
+        for name in ("coef", "domain", "window"):
+            attr = c.pyapi.object_getattr_string(obj, name)
+            with cgutils.early_exit_if_null(c.builder, stack, attr):
+                c.builder.store(cgutils.true_bit, is_error_ptr)
+            t = getattr(typ, name)
+            native = c.unbox(t, attr)
+            c.pyapi.decref(attr)
+            with cgutils.early_exit_if(c.builder, stack, native.is_error):
+                c.builder.store(cgutils.true_bit, is_error_ptr)
+            natives.append(native)
+
+        polynomial.coef = natives[0]
+        polynomial.domain = natives[1]
+        polynomial.window = natives[2]
+
+    return NativeValue(
+        polynomial._getvalue(), is_error=c.builder.load(is_error_ptr)
+    )
+
+
+@box(types.PolynomialType)
+def box_polynomial(typ, val, c):
+    """
+    Convert a native polynomial structure to a Polynomial object.
+    """
+    ret_ptr = cgutils.alloca_once(c.builder, c.pyapi.pyobj)
+    fail_obj = c.pyapi.get_null_object()
+
+    with ExitStack() as stack:
+        polynomial = cgutils.create_struct_proxy(typ)(
+            c.context, c.builder, value=val
+        )
+        coef_obj = c.box(typ.coef, polynomial.coef)
+        with cgutils.early_exit_if_null(c.builder, stack, coef_obj):
+            c.builder.store(fail_obj, ret_ptr)
+
+        domain_obj = c.box(typ.domain, polynomial.domain)
+        with cgutils.early_exit_if_null(c.builder, stack, domain_obj):
+            c.builder.store(fail_obj, ret_ptr)
+
+        window_obj = c.box(typ.window, polynomial.window)
+        with cgutils.early_exit_if_null(c.builder, stack, window_obj):
+            c.builder.store(fail_obj, ret_ptr)
+
+        class_obj = c.pyapi.unserialize(c.pyapi.serialize_object(Polynomial))
+        with cgutils.early_exit_if_null(c.builder, stack, class_obj):
+            c.pyapi.decref(coef_obj)
+            c.pyapi.decref(domain_obj)
+            c.pyapi.decref(window_obj)
+            c.builder.store(fail_obj, ret_ptr)
+
+        if typ.n_args == 1:
+            res1 = c.pyapi.call_function_objargs(class_obj, (coef_obj,))
+            c.builder.store(res1, ret_ptr)
+        else:
+            res3 = c.pyapi.call_function_objargs(
+                class_obj, (coef_obj, domain_obj, window_obj)
+            )
+            c.builder.store(res3, ret_ptr)
+
+        c.pyapi.decref(coef_obj)
+        c.pyapi.decref(domain_obj)
+        c.pyapi.decref(window_obj)
+        c.pyapi.decref(class_obj)
+
+    return c.builder.load(ret_ptr)

numba_cuda/numba/cuda/np/polynomial/polynomial_functions.py

@@ -0,0 +1,380 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+"""
+Implementation of operations involving polynomials.
+"""
+
+import numpy as np
+from numpy.polynomial import polynomial as poly
+from numpy.polynomial import polyutils as pu
+
+from numba.cuda.misc.special import literal_unroll
+from numba.cuda import types
+from numba.cuda.core import errors
+from numba.cuda.extending import overload
+from numba.cuda.np.numpy_support import type_can_asarray, as_dtype, from_dtype
+
+
+@overload(np.roots)
+def roots_impl(p):
+    # cast int vectors to float cf. numpy, this is a bit dicey as
+    # the roots could be complex which will fail anyway
+    ty = getattr(p, "dtype", p)
+    if isinstance(ty, types.Integer):
+        cast_t = np.float64
+    else:
+        cast_t = as_dtype(ty)
+
+    def roots_impl(p):
+        # impl based on numpy:
+        # https://github.com/numpy/numpy/blob/master/numpy/lib/polynomial.py
+
+        if len(p.shape) != 1:
+            raise ValueError("Input must be a 1d array.")
+
+        non_zero = np.nonzero(p)[0]
+
+        if len(non_zero) == 0:
+            return np.zeros(0, dtype=cast_t)
+
+        tz = len(p) - non_zero[-1] - 1
+
+        # pull out the coeffs selecting between possible zero pads
+        p = p[int(non_zero[0]) : int(non_zero[-1]) + 1]
+
+        n = len(p)
+        if n > 1:
+            # construct companion matrix, ensure fortran order
+            # to give to eigvals, write to upper diag and then
+            # transpose.
+            A = np.diag(np.ones((n - 2,), cast_t), 1).T
+            A[0, :] = -p[1:] / p[0]  # normalize
+            roots = np.linalg.eigvals(A)
+        else:
+            roots = np.zeros(0, dtype=cast_t)
+
+        # add in additional zeros on the end if needed
+        if tz > 0:
+            return np.hstack((roots, np.zeros(tz, dtype=cast_t)))
+        else:
+            return roots
+
+    return roots_impl
+
+
+@overload(pu.trimseq)
+def polyutils_trimseq(seq):
+    if not type_can_asarray(seq):
+        msg = 'The argument "seq" must be array-like'
+        raise errors.TypingError(msg)
+
+    if isinstance(seq, types.BaseTuple):
+        msg = 'Unsupported type %r for argument "seq"'
+        raise errors.TypingError(msg % (seq))
+
+    if np.ndim(seq) > 1:
+        msg = "Coefficient array is not 1-d"
+        raise errors.NumbaValueError(msg)
+
+    def impl(seq):
+        if len(seq) == 0:
+            return seq
+        else:
+            for i in range(len(seq) - 1, -1, -1):
+                if seq[i] != 0:
+                    break
+            return seq[: i + 1]
+
+    return impl
+
+
+@overload(pu.as_series)
+def polyutils_as_series(alist, trim=True):
+    if not type_can_asarray(alist):
+        msg = 'The argument "alist" must be array-like'
+        raise errors.TypingError(msg)
+
+    if not isinstance(trim, (bool, types.Boolean)):
+        msg = 'The argument "trim" must be boolean'
+        raise errors.TypingError(msg)
+
+    res_dtype = np.float64
+
+    tuple_input = isinstance(alist, types.BaseTuple)
+    list_input = isinstance(alist, types.List)
+    if tuple_input:
+        if np.any(np.array([np.ndim(a) > 1 for a in alist])):
+            raise errors.NumbaValueError("Coefficient array is not 1-d")
+
+        res_dtype = _poly_result_dtype(*alist)
+
+    elif list_input:
+        dt = as_dtype(_get_list_type(alist))
+        res_dtype = np.result_type(dt, np.float64)
+
+    else:
+        if np.ndim(alist) <= 2:
+            res_dtype = np.result_type(res_dtype, as_dtype(alist.dtype))
+        else:
+            # If total dimension has ndim > 2, then coeff arrays are not 1D
+            raise errors.NumbaValueError("Coefficient array is not 1-d")
+
+    def impl(alist, trim=True):
+        if tuple_input:
+            arrays = []
+            for item in literal_unroll(alist):
+                arrays.append(np.atleast_1d(np.asarray(item)).astype(res_dtype))
+
+        elif list_input:
+            arrays = [
+                np.atleast_1d(np.asarray(a)).astype(res_dtype) for a in alist
+            ]
+
+        else:
+            alist_arr = np.asarray(alist)
+            arrays = [
+                np.atleast_1d(np.asarray(a)).astype(res_dtype)
+                for a in alist_arr
+            ]
+
+        if min([a.size for a in arrays]) == 0:
+            raise ValueError("Coefficient array is empty")
+
+        if trim:
+            arrays = [pu.trimseq(a) for a in arrays]
+
+        ret = arrays
+        return ret
+
+    return impl
+
+
+def _get_list_type(l):
+    # A helper function that takes a list (possibly nested) and returns its
+    # dtype. Returns a Numba type.
+    dt = l.dtype
+    if (not isinstance(dt, types.Number)) and type_can_asarray(dt):
+        return _get_list_type(dt)
+    else:
+        return dt
+
+
+def _poly_result_dtype(*args):
+    # A helper function that takes a tuple of inputs and returns their result
+    # dtype. Used for poly functions. Returns a NumPy dtype.
+    res_dtype = np.float64
+    for item in args:
+        if isinstance(item, types.BaseTuple):
+            s1 = item.types
+        elif isinstance(item, types.List):
+            s1 = [_get_list_type(item)]
+        elif isinstance(item, types.Number):
+            s1 = [item]
+        elif isinstance(item, types.Array):
+            s1 = [item.dtype]
+        else:
+            msg = "Input dtype must be scalar"
+            raise errors.TypingError(msg)
+
+        try:
+            l = [as_dtype(t) for t in s1]
+            l.append(res_dtype)
+            res_dtype = np.result_type(*l)
+        except errors.NumbaNotImplementedError:
+            msg = "Input dtype must be scalar."
+            raise errors.TypingError(msg)
+
+    return from_dtype(res_dtype)
+
+
+@overload(poly.polyadd)
+def numpy_polyadd(c1, c2):
+    if not type_can_asarray(c1):
+        msg = 'The argument "c1" must be array-like'
+        raise errors.TypingError(msg)
+
+    if not type_can_asarray(c2):
+        msg = 'The argument "c2" must be array-like'
+        raise errors.TypingError(msg)
+
+    def impl(c1, c2):
+        arr1, arr2 = pu.as_series((c1, c2))
+        diff = len(arr2) - len(arr1)
+        if diff > 0:
+            zr = np.zeros(diff)
+            arr1 = np.concatenate((arr1, zr))
+        if diff < 0:
+            zr = np.zeros(-diff)
+            arr2 = np.concatenate((arr2, zr))
+        val = arr1 + arr2
+        return pu.trimseq(val)
+
+    return impl
+
+
+@overload(poly.polysub)
+def numpy_polysub(c1, c2):
+    if not type_can_asarray(c1):
+        msg = 'The argument "c1" must be array-like'
+        raise errors.TypingError(msg)
+
+    if not type_can_asarray(c2):
+        msg = 'The argument "c2" must be array-like'
+        raise errors.TypingError(msg)
+
+    def impl(c1, c2):
+        arr1, arr2 = pu.as_series((c1, c2))
+        diff = len(arr2) - len(arr1)
+        if diff > 0:
+            zr = np.zeros(diff)
+            arr1 = np.concatenate((arr1, zr))
+        if diff < 0:
+            zr = np.zeros(-diff)
+            arr2 = np.concatenate((arr2, zr))
+        val = arr1 - arr2
+        return pu.trimseq(val)
+
+    return impl
+
+
+@overload(poly.polymul)
+def numpy_polymul(c1, c2):
+    if not type_can_asarray(c1):
+        msg = 'The argument "c1" must be array-like'
+        raise errors.TypingError(msg)
+
+    if not type_can_asarray(c2):
+        msg = 'The argument "c2" must be array-like'
+        raise errors.TypingError(msg)
+
+    def impl(c1, c2):
+        arr1, arr2 = pu.as_series((c1, c2))
+        val = np.convolve(arr1, arr2)
+        return pu.trimseq(val)
+
+    return impl
+
+
+@overload(poly.polyval, prefer_literal=True)
+def poly_polyval(x, c, tensor=True):
+    if not type_can_asarray(x):
+        msg = 'The argument "x" must be array-like'
+        raise errors.TypingError(msg)
+
+    if not type_can_asarray(c):
+        msg = 'The argument "c" must be array-like'
+        raise errors.TypingError(msg)
+
+    if not isinstance(tensor, (bool, types.BooleanLiteral)):
+        msg = 'The argument "tensor" must be boolean'
+        raise errors.RequireLiteralValue(msg)
+
+    res_dtype = _poly_result_dtype(c, x)
+
+    # Simulate new_shape = (1,) * np.ndim(x) in the general case
+    # If x is a number, new_shape is not used
+    # If x is a tuple or a list, then it's 1d hence new_shape=(1,)
+    x_nd_array = not isinstance(x, types.Number)
+    new_shape = (1,)
+    if isinstance(x, types.Array):
+        # If x is a np.array, then take its dimension
+        new_shape = (1,) * np.ndim(x)
+
+    if isinstance(tensor, bool):
+        tensor_arg = tensor
+    else:
+        tensor_arg = tensor.literal_value
+
+    def impl(x, c, tensor=True):
+        arr = np.asarray(c).astype(res_dtype)
+        inputs = np.asarray(x).astype(res_dtype)
+        if x_nd_array and tensor_arg:
+            arr = arr.reshape(arr.shape + new_shape)
+
+        l = len(arr)
+        y = arr[l - 1] + inputs * 0
+
+        for i in range(l - 1, 0, -1):
+            y = arr[i - 1] + y * inputs
+
+        return y
+
+    return impl
+
+
+@overload(poly.polyint)
+def poly_polyint(c, m=1):
+    if not type_can_asarray(c):
+        msg = 'The argument "c" must be array-like'
+        raise errors.TypingError(msg)
+
+    if not isinstance(m, (int, types.Integer)):
+        msg = 'The argument "m" must be an integer'
+        raise errors.TypingError(msg)
+
+    res_dtype = as_dtype(_poly_result_dtype(c))
+
+    if not np.issubdtype(res_dtype, np.number):
+        msg = f"Input dtype must be scalar. Found {res_dtype} instead"
+        raise errors.TypingError(msg)
+
+    is1D = (np.ndim(c) == 1) or (
+        isinstance(c, (types.List, types.BaseTuple))
+        and isinstance(c.dtype, types.Number)
+    )
+
+    def impl(c, m=1):
+        c = np.asarray(c).astype(res_dtype)
+        cdt = c.dtype
+        for i in range(m):
+            n = len(c)
+
+            tmp = np.empty((n + 1,) + c.shape[1:], dtype=cdt)
+            tmp[0] = c[0] * 0
+            tmp[1] = c[0]
+            for j in range(1, n):
+                tmp[j + 1] = c[j] / (j + 1)
+            c = tmp
+        if is1D:
+            return pu.trimseq(c)
+        else:
+            return c
+
+    return impl
+
+
+@overload(poly.polydiv)
+def numpy_polydiv(c1, c2):
+    if not type_can_asarray(c1):
+        msg = 'The argument "c1" must be array-like'
+        raise errors.TypingError(msg)
+
+    if not type_can_asarray(c2):
+        msg = 'The argument "c2" must be array-like'
+        raise errors.TypingError(msg)
+
+    def impl(c1, c2):
+        arr1, arr2 = pu.as_series((c1, c2))
+        if arr2[-1] == 0:
+            raise ZeroDivisionError()
+
+        l1 = len(arr1)
+        l2 = len(arr2)
+        if l1 < l2:
+            return arr1[:1] * 0, arr1
+        elif l2 == 1:
+            return arr1 / arr2[-1], arr1[:1] * 0
+        else:
+            dlen = l1 - l2
+            scl = arr2[-1]
+            arr2 = arr2[:-1] / scl
+            i = dlen
+            j = l1 - 1
+            while i >= 0:
+                arr1[i:j] -= arr2 * arr1[j]
+                i -= 1
+                j -= 1
+            return arr1[j + 1 :] / scl, pu.trimseq(arr1[: j + 1])
+
+    return impl

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

@@ -0,0 +1,4 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+from numba.cuda.np.ufunc.decorators import vectorize, guvectorize  # noqa: F401