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

numba_cuda/numba/cuda/cpython/rangeobj.py

@@ -0,0 +1,289 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+"""
+Implementation of the range object for fixed-size integers.
+"""
+
+import operator
+
+from numba.cuda import types
+from numba.cuda import cgutils
+from numba.cuda.core.imputils import (
+    Registry,
+    RefType,
+    iternext_impl,
+    impl_ret_untracked,
+)
+from numba.cuda.typing import signature
+from numba.cuda.extending import (
+    intrinsic,
+    overload,
+    overload_attribute,
+    register_jitable,
+)
+
+
+registry = Registry("rangeobj")
+lower = registry.lower
+lower_cast = registry.lower_cast
+
+
+def iterator_impl(iterable_type, iterator_type):
+    """
+    Decorator a given class as implementing *iterator_type*
+    (by providing an `iternext()` method).
+    """
+
+    def wrapper(cls):
+        # These are unbound methods
+        iternext = cls.iternext
+
+        @iternext_impl(RefType.BORROWED)
+        def iternext_wrapper(context, builder, sig, args, result):
+            (value,) = args
+            iterobj = cls(context, builder, value)
+            return iternext(iterobj, context, builder, result)
+
+        lower("iternext", iterator_type)(iternext_wrapper)
+        return cls
+
+    return wrapper
+
+
+def make_range_iterator(typ):
+    """
+    Return the Structure representation of the given *typ* (an
+    instance of types.RangeIteratorType).
+    """
+    return cgutils.create_struct_proxy(typ)
+
+
+def make_range_impl(int_type, range_state_type, range_iter_type):
+    RangeState = cgutils.create_struct_proxy(range_state_type)
+
+    @lower(range, int_type)
+    def range1_impl(context, builder, sig, args):
+        """
+        range(stop: int) -> range object
+        """
+        [stop] = args
+        state = RangeState(context, builder)
+        state.start = context.get_constant(int_type, 0)
+        state.stop = stop
+        state.step = context.get_constant(int_type, 1)
+        return impl_ret_untracked(
+            context, builder, range_state_type, state._getvalue()
+        )
+
+    @lower(range, int_type, int_type)
+    def range2_impl(context, builder, sig, args):
+        """
+        range(start: int, stop: int) -> range object
+        """
+        start, stop = args
+        state = RangeState(context, builder)
+        state.start = start
+        state.stop = stop
+        state.step = context.get_constant(int_type, 1)
+        return impl_ret_untracked(
+            context, builder, range_state_type, state._getvalue()
+        )
+
+    @lower(range, int_type, int_type, int_type)
+    def range3_impl(context, builder, sig, args):
+        """
+        range(start: int, stop: int, step: int) -> range object
+        """
+        [start, stop, step] = args
+        state = RangeState(context, builder)
+        state.start = start
+        state.stop = stop
+        state.step = step
+        return impl_ret_untracked(
+            context, builder, range_state_type, state._getvalue()
+        )
+
+    @lower(len, range_state_type)
+    def range_len(context, builder, sig, args):
+        """
+        len(range)
+        """
+        (value,) = args
+        state = RangeState(context, builder, value)
+        res = RangeIter.from_range_state(context, builder, state)
+        return impl_ret_untracked(
+            context, builder, int_type, builder.load(res.count)
+        )
+
+    @lower("getiter", range_state_type)
+    def getiter_range32_impl(context, builder, sig, args):
+        """
+        range.__iter__
+        """
+        (value,) = args
+        state = RangeState(context, builder, value)
+        res = RangeIter.from_range_state(context, builder, state)._getvalue()
+        return impl_ret_untracked(context, builder, range_iter_type, res)
+
+    @iterator_impl(range_state_type, range_iter_type)
+    class RangeIter(make_range_iterator(range_iter_type)):
+        @classmethod
+        def from_range_state(cls, context, builder, state):
+            """
+            Create a RangeIter initialized from the given RangeState *state*.
+            """
+            self = cls(context, builder)
+            start = state.start
+            stop = state.stop
+            step = state.step
+
+            startptr = cgutils.alloca_once(builder, start.type)
+            builder.store(start, startptr)
+
+            countptr = cgutils.alloca_once(builder, start.type)
+
+            self.iter = startptr
+            self.stop = stop
+            self.step = step
+            self.count = countptr
+
+            diff = builder.sub(stop, start)
+            zero = context.get_constant(int_type, 0)
+            one = context.get_constant(int_type, 1)
+            pos_diff = builder.icmp_signed(">", diff, zero)
+            pos_step = builder.icmp_signed(">", step, zero)
+            sign_differs = builder.xor(pos_diff, pos_step)
+            zero_step = builder.icmp_unsigned("==", step, zero)
+
+            with cgutils.if_unlikely(builder, zero_step):
+                # step shouldn't be zero
+                context.call_conv.return_user_exc(
+                    builder, ValueError, ("range() arg 3 must not be zero",)
+                )
+
+            with builder.if_else(sign_differs) as (then, orelse):
+                with then:
+                    builder.store(zero, self.count)
+
+                with orelse:
+                    rem = builder.srem(diff, step)
+                    rem = builder.select(pos_diff, rem, builder.neg(rem))
+                    uneven = builder.icmp_signed(">", rem, zero)
+                    newcount = builder.add(
+                        builder.sdiv(diff, step),
+                        builder.select(uneven, one, zero),
+                    )
+                    builder.store(newcount, self.count)
+
+            return self
+
+        def iternext(self, context, builder, result):
+            zero = context.get_constant(int_type, 0)
+            countptr = self.count
+            count = builder.load(countptr)
+            is_valid = builder.icmp_signed(">", count, zero)
+            result.set_valid(is_valid)
+
+            with builder.if_then(is_valid):
+                value = builder.load(self.iter)
+                result.yield_(value)
+                one = context.get_constant(int_type, 1)
+
+                builder.store(builder.sub(count, one, flags=["nsw"]), countptr)
+                builder.store(builder.add(value, self.step), self.iter)
+
+
+range_impl_map = {
+    types.int32: (types.range_state32_type, types.range_iter32_type),
+    types.int64: (types.range_state64_type, types.range_iter64_type),
+    types.uint64: (
+        types.unsigned_range_state64_type,
+        types.unsigned_range_iter64_type,
+    ),
+}
+
+for int_type, state_types in range_impl_map.items():
+    make_range_impl(int_type, *state_types)
+
+
+@lower_cast(types.RangeType, types.RangeType)
+def range_to_range(context, builder, fromty, toty, val):
+    olditems = cgutils.unpack_tuple(builder, val, 3)
+    items = [
+        context.cast(builder, v, fromty.dtype, toty.dtype) for v in olditems
+    ]
+    return cgutils.make_anonymous_struct(builder, items)
+
+
+def make_range_attr(index, attribute):
+    @intrinsic
+    def rangetype_attr_getter(typingctx, a):
+        if isinstance(a, types.RangeType):
+
+            def codegen(context, builder, sig, args):
+                (val,) = args
+                items = cgutils.unpack_tuple(builder, val, 3)
+                return impl_ret_untracked(
+                    context, builder, sig.return_type, items[index]
+                )
+
+            return signature(a.dtype, a), codegen
+
+    @overload_attribute(types.RangeType, attribute)
+    def range_attr(rnge):
+        def get(rnge):
+            return rangetype_attr_getter(rnge)
+
+        return get
+
+
+@register_jitable
+def impl_contains_helper(robj, val):
+    if robj.step > 0 and (val < robj.start or val >= robj.stop):
+        return False
+    elif robj.step < 0 and (val <= robj.stop or val > robj.start):
+        return False
+
+    return ((val - robj.start) % robj.step) == 0
+
+
+@overload(operator.contains)
+def impl_contains(robj, val):
+    def impl_false(robj, val):
+        return False
+
+    if not isinstance(robj, types.RangeType):
+        return
+
+    elif isinstance(val, (types.Integer, types.Boolean)):
+        return impl_contains_helper
+
+    elif isinstance(val, types.Float):
+
+        def impl(robj, val):
+            if val % 1 != 0:
+                return False
+            else:
+                return impl_contains_helper(robj, int(val))
+
+        return impl
+
+    elif isinstance(val, types.Complex):
+
+        def impl(robj, val):
+            if val.imag != 0:
+                return False
+            elif val.real % 1 != 0:
+                return False
+            else:
+                return impl_contains_helper(robj, int(val.real))
+
+        return impl
+
+    elif not isinstance(val, types.Number):
+        return impl_false
+
+
+for ix, attr in enumerate(("start", "stop", "step")):
+    make_range_attr(index=ix, attribute=attr)

numba_cuda/numba/cuda/cpython/slicing.py

@@ -0,0 +1,322 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+"""
+Implement slices and various slice computations.
+"""
+
+from llvmlite import ir
+from numba.cuda import types
+from numba.cuda import cgutils
+from numba.cuda.core.imputils import impl_ret_untracked, Registry
+
+registry = Registry("slicing")
+lower = registry.lower
+lower_getattr = registry.lower_getattr
+lower_cast = registry.lower_cast
+lower_constant = registry.lower_constant
+
+
+def fix_index(builder, idx, size):
+    """
+    Fix negative index by adding *size* to it.  Positive
+    indices are left untouched.
+    """
+    is_negative = builder.icmp_signed("<", idx, ir.Constant(size.type, 0))
+    wrapped_index = builder.add(idx, size)
+    return builder.select(is_negative, wrapped_index, idx)
+
+
+def fix_slice(builder, slice, size):
+    """
+    Fix *slice* start and stop to be valid (inclusive and exclusive, resp)
+    indexing bounds for a sequence of the given *size*.
+    """
+    # See PySlice_GetIndicesEx()
+    zero = ir.Constant(size.type, 0)
+    minus_one = ir.Constant(size.type, -1)
+
+    def fix_bound(bound_name, lower_repl, upper_repl):
+        bound = getattr(slice, bound_name)
+        bound = fix_index(builder, bound, size)
+        # Store value
+        setattr(slice, bound_name, bound)
+        # Still negative? => clamp to lower_repl
+        underflow = builder.icmp_signed("<", bound, zero)
+        with builder.if_then(underflow, likely=False):
+            setattr(slice, bound_name, lower_repl)
+        # Greater than size? => clamp to upper_repl
+        overflow = builder.icmp_signed(">=", bound, size)
+        with builder.if_then(overflow, likely=False):
+            setattr(slice, bound_name, upper_repl)
+
+    with builder.if_else(cgutils.is_neg_int(builder, slice.step)) as (
+        if_neg_step,
+        if_pos_step,
+    ):
+        with if_pos_step:
+            # < 0 => 0; >= size => size
+            fix_bound("start", zero, size)
+            fix_bound("stop", zero, size)
+        with if_neg_step:
+            # < 0 => -1; >= size => size - 1
+            lower = minus_one
+            upper = builder.add(size, minus_one)
+            fix_bound("start", lower, upper)
+            fix_bound("stop", lower, upper)
+
+
+def get_slice_length(builder, slicestruct):
+    """
+    Given a slice, compute the number of indices it spans, i.e. the
+    number of iterations that for_range_slice() will execute.
+
+    Pseudo-code:
+        assert step != 0
+        if step > 0:
+            if stop <= start:
+                return 0
+            else:
+                return (stop - start - 1) // step + 1
+        else:
+            if stop >= start:
+                return 0
+            else:
+                return (stop - start + 1) // step + 1
+
+    (see PySlice_GetIndicesEx() in CPython)
+    """
+    start = slicestruct.start
+    stop = slicestruct.stop
+    step = slicestruct.step
+    one = ir.Constant(start.type, 1)
+    zero = ir.Constant(start.type, 0)
+
+    is_step_negative = cgutils.is_neg_int(builder, step)
+    delta = builder.sub(stop, start)
+
+    # Nominal case
+    pos_dividend = builder.sub(delta, one)
+    neg_dividend = builder.add(delta, one)
+    dividend = builder.select(is_step_negative, neg_dividend, pos_dividend)
+    nominal_length = builder.add(one, builder.sdiv(dividend, step))
+
+    # Catch zero length
+    is_zero_length = builder.select(
+        is_step_negative,
+        builder.icmp_signed(">=", delta, zero),
+        builder.icmp_signed("<=", delta, zero),
+    )
+
+    # Clamp to 0 if is_zero_length
+    return builder.select(is_zero_length, zero, nominal_length)
+
+
+def get_slice_bounds(builder, slicestruct):
+    """
+    Return the [lower, upper) indexing bounds of a slice.
+    """
+    start = slicestruct.start
+    stop = slicestruct.stop
+    zero = start.type(0)
+    one = start.type(1)
+    # This is a bit pessimal, e.g. it will return [1, 5) instead
+    # of [1, 4) for `1:5:2`
+    is_step_negative = builder.icmp_signed("<", slicestruct.step, zero)
+    lower = builder.select(is_step_negative, builder.add(stop, one), start)
+    upper = builder.select(is_step_negative, builder.add(start, one), stop)
+    return lower, upper
+
+
+def fix_stride(builder, slice, stride):
+    """
+    Fix the given stride for the slice's step.
+    """
+    return builder.mul(slice.step, stride)
+
+
+def guard_invalid_slice(context, builder, typ, slicestruct):
+    """
+    Guard against *slicestruct* having a zero step (and raise ValueError).
+    """
+    if typ.has_step:
+        cgutils.guard_null(
+            context,
+            builder,
+            slicestruct.step,
+            (ValueError, "slice step cannot be zero"),
+        )
+
+
+def get_defaults(context):
+    """
+    Get the default values for a slice's members:
+    (start for positive step, start for negative step,
+     stop for positive step, stop for negative step, step)
+    """
+    maxint = (1 << (context.address_size - 1)) - 1
+    return (0, maxint, maxint, -maxint - 1, 1)
+
+
+# ---------------------------------------------------------------------------
+# The slice structure
+
+
+@lower(slice, types.VarArg(types.Any))
+def slice_constructor_impl(context, builder, sig, args):
+    (
+        default_start_pos,
+        default_start_neg,
+        default_stop_pos,
+        default_stop_neg,
+        default_step,
+    ) = [context.get_constant(types.intp, x) for x in get_defaults(context)]
+
+    slice_args = [None] * 3
+
+    # Fetch non-None arguments
+    if len(args) == 1 and sig.args[0] is not types.none:
+        slice_args[1] = args[0]
+    else:
+        for i, (ty, val) in enumerate(zip(sig.args, args)):
+            if ty is not types.none:
+                slice_args[i] = val
+
+    # Fill omitted arguments
+    def get_arg_value(i, default):
+        val = slice_args[i]
+        if val is None:
+            return default
+        else:
+            return val
+
+    step = get_arg_value(2, default_step)
+    is_step_negative = builder.icmp_signed(
+        "<", step, context.get_constant(types.intp, 0)
+    )
+    default_stop = builder.select(
+        is_step_negative, default_stop_neg, default_stop_pos
+    )
+    default_start = builder.select(
+        is_step_negative, default_start_neg, default_start_pos
+    )
+    stop = get_arg_value(1, default_stop)
+    start = get_arg_value(0, default_start)
+
+    ty = sig.return_type
+    sli = context.make_helper(builder, sig.return_type)
+    sli.start = start
+    sli.stop = stop
+    sli.step = step
+
+    res = sli._getvalue()
+    return impl_ret_untracked(context, builder, sig.return_type, res)
+
+
+@lower_getattr(types.SliceType, "start")
+def slice_start_impl(context, builder, typ, value):
+    sli = context.make_helper(builder, typ, value)
+    return sli.start
+
+
+@lower_getattr(types.SliceType, "stop")
+def slice_stop_impl(context, builder, typ, value):
+    sli = context.make_helper(builder, typ, value)
+    return sli.stop
+
+
+@lower_getattr(types.SliceType, "step")
+def slice_step_impl(context, builder, typ, value):
+    if typ.has_step:
+        sli = context.make_helper(builder, typ, value)
+        return sli.step
+    else:
+        return context.get_constant(types.intp, 1)
+
+
+@lower("slice.indices", types.SliceType, types.Integer)
+def slice_indices(context, builder, sig, args):
+    length = args[1]
+    sli = context.make_helper(builder, sig.args[0], args[0])
+
+    with builder.if_then(cgutils.is_neg_int(builder, length), likely=False):
+        context.call_conv.return_user_exc(
+            builder, ValueError, ("length should not be negative",)
+        )
+    with builder.if_then(
+        cgutils.is_scalar_zero(builder, sli.step), likely=False
+    ):
+        context.call_conv.return_user_exc(
+            builder, ValueError, ("slice step cannot be zero",)
+        )
+
+    fix_slice(builder, sli, length)
+
+    return context.make_tuple(
+        builder, sig.return_type, (sli.start, sli.stop, sli.step)
+    )
+
+
+def make_slice_from_constant(context, builder, ty, pyval):
+    sli = context.make_helper(builder, ty)
+    lty = context.get_value_type(types.intp)
+
+    (
+        default_start_pos,
+        default_start_neg,
+        default_stop_pos,
+        default_stop_neg,
+        default_step,
+    ) = [context.get_constant(types.intp, x) for x in get_defaults(context)]
+
+    step = pyval.step
+    if step is None:
+        step_is_neg = False
+        step = default_step
+    else:
+        step_is_neg = step < 0
+        step = lty(step)
+
+    start = pyval.start
+    if start is None:
+        if step_is_neg:
+            start = default_start_neg
+        else:
+            start = default_start_pos
+    else:
+        start = lty(start)
+
+    stop = pyval.stop
+    if stop is None:
+        if step_is_neg:
+            stop = default_stop_neg
+        else:
+            stop = default_stop_pos
+    else:
+        stop = lty(stop)
+
+    sli.start = start
+    sli.stop = stop
+    sli.step = step
+
+    return sli._getvalue()
+
+
+@lower_constant(types.SliceType)
+def constant_slice(context, builder, ty, pyval):
+    if isinstance(ty, types.Literal):
+        typ = ty.literal_type
+    else:
+        typ = ty
+
+    return make_slice_from_constant(context, builder, typ, pyval)
+
+
+@lower_cast(types.misc.SliceLiteral, types.SliceType)
+def cast_from_literal(context, builder, fromty, toty, val):
+    return make_slice_from_constant(
+        context,
+        builder,
+        toty,
+        fromty.literal_value,
+    )