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

numba_cuda/numba/cuda/np/unsafe/ndarray.py

@@ -0,0 +1,84 @@
+# 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 numpy.
+"""
+
+from numba.cuda import types
+from numba.cuda.cgutils import unpack_tuple
+from numba.cuda.extending import intrinsic
+from numba.cuda import typing
+from numba.cuda.core.imputils import impl_ret_new_ref
+from numba.cuda.core.errors import RequireLiteralValue, TypingError
+
+from numba.cuda.cpython.unsafe.tuple import tuple_setitem
+
+
+@intrinsic
+def empty_inferred(typingctx, shape):
+    """A version of numpy.empty whose dtype is inferred by the type system.
+
+    Expects `shape` to be a int-tuple.
+
+    There is special logic in the type-inferencer to handle the "refine"-ing
+    of undefined dtype.
+    """
+    from numba.cuda.np.arrayobj import _empty_nd_impl
+
+    def codegen(context, builder, signature, args):
+        # check that the return type is now defined
+        arrty = signature.return_type
+        assert arrty.is_precise()
+        shapes = unpack_tuple(builder, args[0])
+        # redirect implementation to np.empty
+        res = _empty_nd_impl(context, builder, arrty, shapes)
+        return impl_ret_new_ref(context, builder, arrty, res._getvalue())
+
+    # make function signature
+    nd = len(shape)
+    array_ty = types.Array(ndim=nd, layout="C", dtype=types.undefined)
+    sig = array_ty(shape)
+    return sig, codegen
+
+
+@intrinsic
+def to_fixed_tuple(typingctx, array, length):
+    """Convert *array* into a tuple of *length*
+
+    Returns ``UniTuple(array.dtype, length)``
+
+    ** Warning **
+    - No boundchecking.
+      If *length* is longer than *array.size*, the behavior is undefined.
+    """
+    if not isinstance(length, types.IntegerLiteral):
+        raise RequireLiteralValue("*length* argument must be a constant")
+
+    if array.ndim != 1:
+        raise TypingError("Not supported on array.ndim={}".format(array.ndim))
+
+    # Determine types
+    tuple_size = int(length.literal_value)
+    tuple_type = types.UniTuple(dtype=array.dtype, count=tuple_size)
+    sig = tuple_type(array, length)
+
+    def codegen(context, builder, signature, args):
+        def impl(array, length, empty_tuple):
+            out = empty_tuple
+            for i in range(length):
+                out = tuple_setitem(out, i, array[i])
+            return out
+
+        inner_argtypes = [signature.args[0], 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 = [args[0], ll_idx_type(tuple_size), empty_tuple]
+
+        res = context.compile_internal(builder, impl, inner_sig, inner_args)
+        return res
+
+    return sig, codegen

numba_cuda/numba/cuda/nvvmutils.py

@@ -0,0 +1,254 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import itertools
+from llvmlite import ir
+from numba.cuda.core import targetconfig
+from numba.cuda import cgutils
+from .cudadrv import nvvm
+
+
+def declare_atomic_cas_int(lmod, isize):
+    fname = "___numba_atomic_i" + str(isize) + "_cas_hack"
+    fnty = ir.FunctionType(
+        ir.IntType(isize),
+        (
+            ir.PointerType(ir.IntType(isize)),
+            ir.IntType(isize),
+            ir.IntType(isize),
+        ),
+    )
+    return cgutils.get_or_insert_function(lmod, fnty, fname)
+
+
+def atomic_cmpxchg(builder, lmod, isize, ptr, cmp, val):
+    out = builder.cmpxchg(ptr, cmp, val, "monotonic", "monotonic")
+    return builder.extract_value(out, 0)
+
+
+def declare_atomic_add_float32(lmod):
+    fname = "llvm.nvvm.atomic.load.add.f32.p0f32"
+    fnty = ir.FunctionType(
+        ir.FloatType(), (ir.PointerType(ir.FloatType(), 0), ir.FloatType())
+    )
+    return cgutils.get_or_insert_function(lmod, fnty, fname)
+
+
+def declare_atomic_add_float64(lmod):
+    flags = targetconfig.ConfigStack().top()
+    if flags.compute_capability >= (6, 0):
+        fname = "llvm.nvvm.atomic.load.add.f64.p0f64"
+    else:
+        fname = "___numba_atomic_double_add"
+    fnty = ir.FunctionType(
+        ir.DoubleType(), (ir.PointerType(ir.DoubleType()), ir.DoubleType())
+    )
+    return cgutils.get_or_insert_function(lmod, fnty, fname)
+
+
+def declare_atomic_sub_float32(lmod):
+    fname = "___numba_atomic_float_sub"
+    fnty = ir.FunctionType(
+        ir.FloatType(), (ir.PointerType(ir.FloatType()), ir.FloatType())
+    )
+    return cgutils.get_or_insert_function(lmod, fnty, fname)
+
+
+def declare_atomic_sub_float64(lmod):
+    fname = "___numba_atomic_double_sub"
+    fnty = ir.FunctionType(
+        ir.DoubleType(), (ir.PointerType(ir.DoubleType()), ir.DoubleType())
+    )
+    return cgutils.get_or_insert_function(lmod, fnty, fname)
+
+
+def declare_atomic_inc_int32(lmod):
+    fname = "llvm.nvvm.atomic.load.inc.32.p0i32"
+    fnty = ir.FunctionType(
+        ir.IntType(32), (ir.PointerType(ir.IntType(32)), ir.IntType(32))
+    )
+    return cgutils.get_or_insert_function(lmod, fnty, fname)
+
+
+def declare_atomic_inc_int64(lmod):
+    fname = "___numba_atomic_u64_inc"
+    fnty = ir.FunctionType(
+        ir.IntType(64), (ir.PointerType(ir.IntType(64)), ir.IntType(64))
+    )
+    return cgutils.get_or_insert_function(lmod, fnty, fname)
+
+
+def declare_atomic_dec_int32(lmod):
+    fname = "llvm.nvvm.atomic.load.dec.32.p0i32"
+    fnty = ir.FunctionType(
+        ir.IntType(32), (ir.PointerType(ir.IntType(32)), ir.IntType(32))
+    )
+    return cgutils.get_or_insert_function(lmod, fnty, fname)
+
+
+def declare_atomic_dec_int64(lmod):
+    fname = "___numba_atomic_u64_dec"
+    fnty = ir.FunctionType(
+        ir.IntType(64), (ir.PointerType(ir.IntType(64)), ir.IntType(64))
+    )
+    return cgutils.get_or_insert_function(lmod, fnty, fname)
+
+
+def declare_atomic_max_float32(lmod):
+    fname = "___numba_atomic_float_max"
+    fnty = ir.FunctionType(
+        ir.FloatType(), (ir.PointerType(ir.FloatType()), ir.FloatType())
+    )
+    return cgutils.get_or_insert_function(lmod, fnty, fname)
+
+
+def declare_atomic_max_float64(lmod):
+    fname = "___numba_atomic_double_max"
+    fnty = ir.FunctionType(
+        ir.DoubleType(), (ir.PointerType(ir.DoubleType()), ir.DoubleType())
+    )
+    return cgutils.get_or_insert_function(lmod, fnty, fname)
+
+
+def declare_atomic_min_float32(lmod):
+    fname = "___numba_atomic_float_min"
+    fnty = ir.FunctionType(
+        ir.FloatType(), (ir.PointerType(ir.FloatType()), ir.FloatType())
+    )
+    return cgutils.get_or_insert_function(lmod, fnty, fname)
+
+
+def declare_atomic_min_float64(lmod):
+    fname = "___numba_atomic_double_min"
+    fnty = ir.FunctionType(
+        ir.DoubleType(), (ir.PointerType(ir.DoubleType()), ir.DoubleType())
+    )
+    return cgutils.get_or_insert_function(lmod, fnty, fname)
+
+
+def declare_atomic_nanmax_float32(lmod):
+    fname = "___numba_atomic_float_nanmax"
+    fnty = ir.FunctionType(
+        ir.FloatType(), (ir.PointerType(ir.FloatType()), ir.FloatType())
+    )
+    return cgutils.get_or_insert_function(lmod, fnty, fname)
+
+
+def declare_atomic_nanmax_float64(lmod):
+    fname = "___numba_atomic_double_nanmax"
+    fnty = ir.FunctionType(
+        ir.DoubleType(), (ir.PointerType(ir.DoubleType()), ir.DoubleType())
+    )
+    return cgutils.get_or_insert_function(lmod, fnty, fname)
+
+
+def declare_atomic_nanmin_float32(lmod):
+    fname = "___numba_atomic_float_nanmin"
+    fnty = ir.FunctionType(
+        ir.FloatType(), (ir.PointerType(ir.FloatType()), ir.FloatType())
+    )
+    return cgutils.get_or_insert_function(lmod, fnty, fname)
+
+
+def declare_atomic_nanmin_float64(lmod):
+    fname = "___numba_atomic_double_nanmin"
+    fnty = ir.FunctionType(
+        ir.DoubleType(), (ir.PointerType(ir.DoubleType()), ir.DoubleType())
+    )
+    return cgutils.get_or_insert_function(lmod, fnty, fname)
+
+
+def declare_cudaCGGetIntrinsicHandle(lmod):
+    fname = "cudaCGGetIntrinsicHandle"
+    fnty = ir.FunctionType(ir.IntType(64), (ir.IntType(32),))
+    return cgutils.get_or_insert_function(lmod, fnty, fname)
+
+
+def declare_cudaCGSynchronize(lmod):
+    fname = "cudaCGSynchronize"
+    fnty = ir.FunctionType(ir.IntType(32), (ir.IntType(64), ir.IntType(32)))
+    return cgutils.get_or_insert_function(lmod, fnty, fname)
+
+
+def declare_string(builder, value):
+    lmod = builder.basic_block.function.module
+    cval = cgutils.make_bytearray(value.encode("utf-8") + b"\x00")
+    gl = cgutils.add_global_variable(
+        lmod, cval.type, name="_str", addrspace=nvvm.ADDRSPACE_CONSTANT
+    )
+    gl.linkage = "internal"
+    gl.global_constant = True
+    gl.initializer = cval
+
+    return builder.addrspacecast(gl, ir.PointerType(ir.IntType(8)), "generic")
+
+
+def declare_vprint(lmod):
+    voidptrty = ir.PointerType(ir.IntType(8))
+    # NOTE: the second argument to vprintf() points to the variable-length
+    # array of arguments (after the format)
+    vprintfty = ir.FunctionType(ir.IntType(32), [voidptrty, voidptrty])
+    vprintf = cgutils.get_or_insert_function(lmod, vprintfty, "vprintf")
+    return vprintf
+
+
+# -----------------------------------------------------------------------------
+
+SREG_MAPPING = {
+    "tid.x": "llvm.nvvm.read.ptx.sreg.tid.x",
+    "tid.y": "llvm.nvvm.read.ptx.sreg.tid.y",
+    "tid.z": "llvm.nvvm.read.ptx.sreg.tid.z",
+    "ntid.x": "llvm.nvvm.read.ptx.sreg.ntid.x",
+    "ntid.y": "llvm.nvvm.read.ptx.sreg.ntid.y",
+    "ntid.z": "llvm.nvvm.read.ptx.sreg.ntid.z",
+    "ctaid.x": "llvm.nvvm.read.ptx.sreg.ctaid.x",
+    "ctaid.y": "llvm.nvvm.read.ptx.sreg.ctaid.y",
+    "ctaid.z": "llvm.nvvm.read.ptx.sreg.ctaid.z",
+    "nctaid.x": "llvm.nvvm.read.ptx.sreg.nctaid.x",
+    "nctaid.y": "llvm.nvvm.read.ptx.sreg.nctaid.y",
+    "nctaid.z": "llvm.nvvm.read.ptx.sreg.nctaid.z",
+    "warpsize": "llvm.nvvm.read.ptx.sreg.warpsize",
+    "laneid": "llvm.nvvm.read.ptx.sreg.laneid",
+}
+
+
+def call_sreg(builder, name):
+    module = builder.module
+    fnty = ir.FunctionType(ir.IntType(32), ())
+    fn = cgutils.get_or_insert_function(module, fnty, SREG_MAPPING[name])
+    return builder.call(fn, ())
+
+
+class SRegBuilder(object):
+    def __init__(self, builder):
+        self.builder = builder
+
+    def tid(self, xyz):
+        return call_sreg(self.builder, "tid.%s" % xyz)
+
+    def ctaid(self, xyz):
+        return call_sreg(self.builder, "ctaid.%s" % xyz)
+
+    def ntid(self, xyz):
+        return call_sreg(self.builder, "ntid.%s" % xyz)
+
+    def nctaid(self, xyz):
+        return call_sreg(self.builder, "nctaid.%s" % xyz)
+
+    def getdim(self, xyz):
+        i64 = ir.IntType(64)
+        tid = self.builder.sext(self.tid(xyz), i64)
+        ntid = self.builder.sext(self.ntid(xyz), i64)
+        nctaid = self.builder.sext(self.ctaid(xyz), i64)
+        res = self.builder.add(self.builder.mul(ntid, nctaid), tid)
+        return res
+
+
+def get_global_id(builder, dim):
+    sreg = SRegBuilder(builder)
+    it = (sreg.getdim(xyz) for xyz in "xyz")
+    seq = list(itertools.islice(it, None, dim))
+    if dim == 1:
+        return seq[0]
+    else:
+        return seq

numba_cuda/numba/cuda/printimpl.py

@@ -0,0 +1,126 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+from functools import singledispatch
+from llvmlite import ir
+from numba.cuda import types
+from numba.cuda import cgutils
+from numba.cuda.core.errors import NumbaWarning
+from numba.cuda.core.imputils import Registry
+from numba.cuda import nvvmutils
+from numba.cuda.types.ext_types import Dim3, Bfloat16
+from warnings import warn
+
+registry = Registry("printimpl")
+lower = registry.lower
+
+voidptr = ir.PointerType(ir.IntType(8))
+
+
+# NOTE: we don't use @lower here since print_item() doesn't return a LLVM value
+
+
+@singledispatch
+def print_item(ty, context, builder, val):
+    """
+    Handle printing of a single value of the given Numba type.
+    A (format string, [list of arguments]) is returned that will allow
+    forming the final printf()-like call.
+    """
+    raise NotImplementedError(
+        "printing unimplemented for values of type %s" % (ty,)
+    )
+
+
+@print_item.register(types.Integer)
+@print_item.register(types.IntegerLiteral)
+def int_print_impl(ty, context, builder, val):
+    if ty in types.unsigned_domain:
+        rawfmt = "%llu"
+        dsttype = types.uint64
+    else:
+        rawfmt = "%lld"
+        dsttype = types.int64
+    lld = context.cast(builder, val, ty, dsttype)
+    return rawfmt, [lld]
+
+
+@print_item.register(types.Float)
+def real_print_impl(ty, context, builder, val):
+    lld = context.cast(builder, val, ty, types.float64)
+    return "%f", [lld]
+
+
+@print_item.register(Bfloat16)
+def bfloat16_print_impl(ty, context, builder, val):
+    # Hand rolled bfloat16 -> float32 -> double conversion with zero-ext
+    bits32 = builder.zext(val, ir.IntType(32))
+    shift = builder.shl(bits32, ir.Constant(ir.IntType(32), 16))
+    f32 = builder.bitcast(shift, ir.FloatType())
+    # printf("%f") expects a double; promote to f64 to match vararg expectation
+    f64 = builder.fpext(f32, ir.DoubleType())
+    return "%f", [f64]
+
+
+@print_item.register(types.StringLiteral)
+def const_print_impl(ty, context, builder, sigval):
+    pyval = ty.literal_value
+    assert isinstance(pyval, str)  # Ensured by lowering
+    rawfmt = "%s"
+    val = context.insert_string_const_addrspace(builder, pyval)
+    return rawfmt, [val]
+
+
+@print_item.register(Dim3)
+def dim3_print_impl(ty, context, builder, val):
+    rawfmt = "(%d, %d, %d)"
+    x = builder.extract_value(val, 0)
+    y = builder.extract_value(val, 1)
+    z = builder.extract_value(val, 2)
+    return rawfmt, [x, y, z]
+
+
+@print_item.register(types.Boolean)
+def bool_print_impl(ty, context, builder, val):
+    true_string = context.insert_string_const_addrspace(builder, "True")
+    false_string = context.insert_string_const_addrspace(builder, "False")
+    res_ptr = cgutils.alloca_once_value(builder, false_string)
+    with builder.if_then(val):
+        builder.store(true_string, res_ptr)
+    rawfmt = "%s"
+    return rawfmt, [builder.load(res_ptr)]
+
+
+@lower(print, types.VarArg(types.Any))
+def print_varargs(context, builder, sig, args):
+    """This function is a generic 'print' wrapper for arbitrary types.
+    It dispatches to the appropriate 'print' implementations above
+    depending on the detected real types in the signature."""
+
+    vprint = nvvmutils.declare_vprint(builder.module)
+
+    formats = []
+    values = []
+
+    for i, (argtype, argval) in enumerate(zip(sig.args, args)):
+        argfmt, argvals = print_item(argtype, context, builder, argval)
+        formats.append(argfmt)
+        values.extend(argvals)
+
+    rawfmt = " ".join(formats) + "\n"
+    if len(args) > 32:
+        msg = (
+            "CUDA print() cannot print more than 32 items. "
+            "The raw format string will be emitted by the kernel instead."
+        )
+        warn(msg, NumbaWarning)
+
+        rawfmt = rawfmt.replace("%", "%%")
+    fmt = context.insert_string_const_addrspace(builder, rawfmt)
+    array = cgutils.make_anonymous_struct(builder, values)
+    arrayptr = cgutils.alloca_once_value(builder, array)
+
+    vprint = nvvmutils.declare_vprint(builder.module)
+    builder.call(vprint, (fmt, builder.bitcast(arrayptr, voidptr)))
+
+    return context.get_dummy_value()