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

numba_cuda/numba/cuda/cpython/unsafe/numbers.py

@@ -0,0 +1,64 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+"""This module provides the unsafe things for targets/numbers.py"""
+
+from numba.cuda import types
+from numba.cuda.core import errors
+from numba.cuda.extending import intrinsic
+
+from llvmlite import ir
+
+
+@intrinsic
+def viewer(tyctx, val, viewty):
+    """Bitcast a scalar 'val' to the given type 'viewty'."""
+    bits = val.bitwidth
+    if isinstance(viewty.dtype, types.Integer):
+        bitcastty = ir.IntType(bits)
+    elif isinstance(viewty.dtype, types.Float):
+        bitcastty = ir.FloatType() if bits == 32 else ir.DoubleType()
+    else:
+        assert 0, "unreachable"
+
+    def codegen(cgctx, builder, typ, args):
+        flt = args[0]
+        return builder.bitcast(flt, bitcastty)
+
+    retty = viewty.dtype
+    sig = retty(val, viewty)
+    return sig, codegen
+
+
+@intrinsic
+def trailing_zeros(typeingctx, src):
+    """Counts trailing zeros in the binary representation of an integer."""
+    if not isinstance(src, types.Integer):
+        msg = (
+            "trailing_zeros is only defined for integers, but value passed "
+            f"was '{src}'."
+        )
+        raise errors.NumbaTypeError(msg)
+
+    def codegen(context, builder, signature, args):
+        [src] = args
+        return builder.cttz(src, ir.Constant(ir.IntType(1), 0))
+
+    return src(src), codegen
+
+
+@intrinsic
+def leading_zeros(typeingctx, src):
+    """Counts leading zeros in the binary representation of an integer."""
+    if not isinstance(src, types.Integer):
+        msg = (
+            "leading_zeros is only defined for integers, but value passed "
+            f"was '{src}'."
+        )
+        raise errors.NumbaTypeError(msg)
+
+    def codegen(context, builder, signature, args):
+        [src] = args
+        return builder.ctlz(src, ir.Constant(ir.IntType(1), 0))
+
+    return src(src), codegen

numba_cuda/numba/cuda/cpython/unsafe/tuple.py

@@ -0,0 +1,92 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+"""
+This file provides internal compiler utilities that support certain special
+operations with tuple and workarounds for limitations enforced in userland.
+"""
+
+from numba.cuda import types, typing
+from numba.cuda.core import errors
+from numba.cuda.cgutils import alloca_once
+from numba.cuda.extending import intrinsic
+
+
+@intrinsic
+def tuple_setitem(typingctx, tup, idx, val):
+    """Return a copy of the tuple with item at *idx* replaced with *val*.
+
+    Operation: ``out = tup[:idx] + (val,) + tup[idx + 1:]
+
+    **Warning**
+
+    - No boundchecking.
+    - The dtype of the tuple cannot be changed.
+      *val* is always cast to the existing dtype of the tuple.
+    """
+
+    def codegen(context, builder, signature, args):
+        tup, idx, val = args
+        stack = alloca_once(builder, tup.type)
+        builder.store(tup, stack)
+        # Unsafe load on unchecked bounds.  Poison value maybe returned.
+        offptr = builder.gep(stack, [idx.type(0), idx], inbounds=True)
+        builder.store(val, offptr)
+        return builder.load(stack)
+
+    sig = tup(tup, idx, tup.dtype)
+    return sig, codegen
+
+
+@intrinsic
+def build_full_slice_tuple(tyctx, sz):
+    """Creates a sz-tuple of full slices."""
+    if not isinstance(sz, types.IntegerLiteral):
+        raise errors.RequireLiteralValue(sz)
+
+    size = int(sz.literal_value)
+    tuple_type = types.UniTuple(dtype=types.slice2_type, count=size)
+    sig = tuple_type(sz)
+
+    def codegen(context, builder, signature, args):
+        def impl(length, empty_tuple):
+            out = empty_tuple
+            for i in range(length):
+                out = tuple_setitem(out, i, slice(None, None))
+            return out
+
+        inner_argtypes = [types.intp, tuple_type]
+        inner_sig = typing.signature(tuple_type, *inner_argtypes)
+        ll_idx_type = context.get_value_type(types.intp)
+        # Allocate an empty tuple
+        empty_tuple = context.get_constant_undef(tuple_type)
+        inner_args = [ll_idx_type(size), empty_tuple]
+
+        res = context.compile_internal(builder, impl, inner_sig, inner_args)
+        return res
+
+    return sig, codegen
+
+
+@intrinsic
+def unpack_single_tuple(tyctx, tup):
+    """This exists to handle the situation y = (*x,), the interpreter injects a
+    call to it in the case of a single value unpack. It's not possible at
+    interpreting time to differentiate between an unpack on a variable sized
+    container e.g. list and a fixed one, e.g. tuple. This function handles the
+    situation should it arise.
+    """
+    # See issue #6534
+    if not isinstance(tup, types.BaseTuple):
+        msg = (
+            f"Only tuples are supported when unpacking a single item, "
+            f"got type: {tup}"
+        )
+        raise errors.UnsupportedError(msg)
+
+    sig = tup(tup)
+
+    def codegen(context, builder, signature, args):
+        return args[0]  # there's only one tuple and it's a simple pass through
+
+    return sig, codegen