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

numba_cuda/numba/cuda/cudaimpl.py

@@ -0,0 +1,983 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+from functools import reduce
+import operator
+import math
+import struct
+
+from llvmlite import ir
+import llvmlite.binding as ll
+
+from numba.cuda.core.imputils import Registry
+from numba.cuda.typing.npydecl import parse_dtype
+from numba.cuda.datamodel.models import StructModel
+from numba.cuda import types
+from numba.cuda import cgutils
+from numba.cuda.np import ufunc_db
+from numba.cuda.np.npyimpl import register_ufuncs
+from .cudadrv import nvvm
+from numba import cuda
+from numba.cuda.api_util import normalize_indices
+from numba.cuda import nvvmutils, stubs
+from numba.cuda.types.ext_types import dim3, CUDADispatcher
+
+if cuda.HAS_NUMBA:
+    from numba.core.datamodel.models import StructModel as CoreStructModel
+    from numba.core import types as core_types
+
+registry = Registry("cudaimpl")
+lower = registry.lower
+lower_attr = registry.lower_getattr
+lower_constant = registry.lower_constant
+lower_getattr_generic = registry.lower_getattr_generic
+lower_setattr = registry.lower_setattr
+lower_setattr_generic = registry.lower_setattr_generic
+lower_cast = registry.lower_cast
+
+
+def initialize_dim3(builder, prefix):
+    x = nvvmutils.call_sreg(builder, "%s.x" % prefix)
+    y = nvvmutils.call_sreg(builder, "%s.y" % prefix)
+    z = nvvmutils.call_sreg(builder, "%s.z" % prefix)
+    return cgutils.pack_struct(builder, (x, y, z))
+
+
+@lower_attr(types.Module(cuda), "threadIdx")
+def cuda_threadIdx(context, builder, sig, args):
+    return initialize_dim3(builder, "tid")
+
+
+@lower_attr(types.Module(cuda), "blockDim")
+def cuda_blockDim(context, builder, sig, args):
+    return initialize_dim3(builder, "ntid")
+
+
+@lower_attr(types.Module(cuda), "blockIdx")
+def cuda_blockIdx(context, builder, sig, args):
+    return initialize_dim3(builder, "ctaid")
+
+
+@lower_attr(types.Module(cuda), "gridDim")
+def cuda_gridDim(context, builder, sig, args):
+    return initialize_dim3(builder, "nctaid")
+
+
+@lower_attr(types.Module(cuda), "laneid")
+def cuda_laneid(context, builder, sig, args):
+    return nvvmutils.call_sreg(builder, "laneid")
+
+
+@lower_attr(dim3, "x")
+def dim3_x(context, builder, sig, args):
+    return builder.extract_value(args, 0)
+
+
+@lower_attr(dim3, "y")
+def dim3_y(context, builder, sig, args):
+    return builder.extract_value(args, 1)
+
+
+@lower_attr(dim3, "z")
+def dim3_z(context, builder, sig, args):
+    return builder.extract_value(args, 2)
+
+
+# -----------------------------------------------------------------------------
+
+
+@lower(cuda.const.array_like, types.Array)
+def cuda_const_array_like(context, builder, sig, args):
+    # This is a no-op because CUDATargetContext.make_constant_array already
+    # created the constant array.
+    return args[0]
+
+
+_unique_smem_id = 0
+
+
+def _get_unique_smem_id(name):
+    """Due to bug with NVVM invalid internalizing of shared memory in the
+    PTX output.  We can't mark shared memory to be internal. We have to
+    ensure unique name is generated for shared memory symbol.
+    """
+    global _unique_smem_id
+    _unique_smem_id += 1
+    return "{0}_{1}".format(name, _unique_smem_id)
+
+
+def _validate_alignment(alignment: int):
+    """
+    Ensures that *alignment*, if not None, is a) greater than zero, b) a power
+    of two, and c) a multiple of the size of a pointer.  If any of these
+    conditions are not met, a ValueError is raised.  Otherwise, this
+    function returns None, indicating that the alignment is valid.
+    """
+    if alignment is None:
+        return
+    if not isinstance(alignment, int):
+        raise ValueError("Alignment must be an integer")
+    if alignment <= 0:
+        raise ValueError("Alignment must be positive")
+    if (alignment & (alignment - 1)) != 0:
+        raise ValueError("Alignment must be a power of 2")
+    pointer_size = struct.calcsize("P")
+    if (alignment % pointer_size) != 0:
+        msg = f"Alignment must be a multiple of {pointer_size}"
+        raise ValueError(msg)
+
+
+def _try_extract_and_validate_alignment(sig: types.Tuple):
+    """
+    Extracts and validates the alignment from the supplied signature.
+
+    Returns the alignment if it is present and is an integer literal;
+    otherwise, returns None.
+
+    N.B. Currently, this routine assumes the signature has exactly
+         three arguments, with the alignment (if present) as the third
+         argument, as is the case with the shared and local array
+         helper routines below.
+
+         If this routine is called from new places, you may need to
+         review this implicit assumption.
+    """
+    if len(sig.args) != 3:
+        return None
+
+    alignment_arg = sig.args[2]
+    if not isinstance(alignment_arg, types.IntegerLiteral):
+        return None
+
+    alignment_arg = alignment_arg.literal_value
+    _validate_alignment(alignment_arg)
+    return alignment_arg
+
+
+@lower(cuda.shared.array, types.IntegerLiteral, types.Any)
+@lower(cuda.shared.array, types.IntegerLiteral, types.Any, types.IntegerLiteral)
+@lower(cuda.shared.array, types.IntegerLiteral, types.Any, types.NoneType)
+def cuda_shared_array_integer(context, builder, sig, args):
+    length = sig.args[0].literal_value
+    dtype = parse_dtype(sig.args[1])
+    alignment = _try_extract_and_validate_alignment(sig)
+    return _generic_array(
+        context,
+        builder,
+        shape=(length,),
+        dtype=dtype,
+        symbol_name=_get_unique_smem_id("_cudapy_smem"),
+        addrspace=nvvm.ADDRSPACE_SHARED,
+        can_dynsized=True,
+        alignment=alignment,
+    )
+
+
+@lower(cuda.shared.array, types.BaseTuple, types.Any)
+@lower(cuda.shared.array, types.BaseTuple, types.Any, types.IntegerLiteral)
+@lower(cuda.shared.array, types.BaseTuple, types.Any, types.NoneType)
+def cuda_shared_array_tuple(context, builder, sig, args):
+    shape = [s.literal_value for s in sig.args[0]]
+    dtype = parse_dtype(sig.args[1])
+    alignment = _try_extract_and_validate_alignment(sig)
+    return _generic_array(
+        context,
+        builder,
+        shape=shape,
+        dtype=dtype,
+        symbol_name=_get_unique_smem_id("_cudapy_smem"),
+        addrspace=nvvm.ADDRSPACE_SHARED,
+        can_dynsized=True,
+        alignment=alignment,
+    )
+
+
+@lower(cuda.local.array, types.IntegerLiteral, types.Any)
+@lower(cuda.local.array, types.IntegerLiteral, types.Any, types.IntegerLiteral)
+@lower(cuda.local.array, types.IntegerLiteral, types.Any, types.NoneType)
+def cuda_local_array_integer(context, builder, sig, args):
+    length = sig.args[0].literal_value
+    dtype = parse_dtype(sig.args[1])
+    alignment = _try_extract_and_validate_alignment(sig)
+    return _generic_array(
+        context,
+        builder,
+        shape=(length,),
+        dtype=dtype,
+        symbol_name="_cudapy_lmem",
+        addrspace=nvvm.ADDRSPACE_LOCAL,
+        can_dynsized=False,
+        alignment=alignment,
+    )
+
+
+@lower(cuda.local.array, types.BaseTuple, types.Any)
+@lower(cuda.local.array, types.BaseTuple, types.Any, types.IntegerLiteral)
+@lower(cuda.local.array, types.BaseTuple, types.Any, types.NoneType)
+def cuda_local_array_tuple(context, builder, sig, args):
+    shape = [s.literal_value for s in sig.args[0]]
+    dtype = parse_dtype(sig.args[1])
+    alignment = _try_extract_and_validate_alignment(sig)
+    return _generic_array(
+        context,
+        builder,
+        shape=shape,
+        dtype=dtype,
+        symbol_name="_cudapy_lmem",
+        addrspace=nvvm.ADDRSPACE_LOCAL,
+        can_dynsized=False,
+        alignment=alignment,
+    )
+
+
+@lower(stubs.threadfence_block)
+def ptx_threadfence_block(context, builder, sig, args):
+    assert not args
+    fname = "llvm.nvvm.membar.cta"
+    lmod = builder.module
+    fnty = ir.FunctionType(ir.VoidType(), ())
+    sync = cgutils.get_or_insert_function(lmod, fnty, fname)
+    builder.call(sync, ())
+    return context.get_dummy_value()
+
+
+@lower(stubs.threadfence_system)
+def ptx_threadfence_system(context, builder, sig, args):
+    assert not args
+    fname = "llvm.nvvm.membar.sys"
+    lmod = builder.module
+    fnty = ir.FunctionType(ir.VoidType(), ())
+    sync = cgutils.get_or_insert_function(lmod, fnty, fname)
+    builder.call(sync, ())
+    return context.get_dummy_value()
+
+
+@lower(stubs.threadfence)
+def ptx_threadfence_device(context, builder, sig, args):
+    assert not args
+    fname = "llvm.nvvm.membar.gl"
+    lmod = builder.module
+    fnty = ir.FunctionType(ir.VoidType(), ())
+    sync = cgutils.get_or_insert_function(lmod, fnty, fname)
+    builder.call(sync, ())
+    return context.get_dummy_value()
+
+
+@lower(stubs.syncwarp)
+def ptx_syncwarp(context, builder, sig, args):
+    mask = context.get_constant(types.int32, 0xFFFFFFFF)
+    mask_sig = types.none(types.int32)
+    return ptx_syncwarp_mask(context, builder, mask_sig, [mask])
+
+
+@lower(stubs.syncwarp, types.i4)
+def ptx_syncwarp_mask(context, builder, sig, args):
+    fname = "llvm.nvvm.bar.warp.sync"
+    lmod = builder.module
+    fnty = ir.FunctionType(ir.VoidType(), (ir.IntType(32),))
+    sync = cgutils.get_or_insert_function(lmod, fnty, fname)
+    builder.call(sync, args)
+    return context.get_dummy_value()
+
+
+@lower(stubs.match_any_sync, types.i4, types.i4)
+@lower(stubs.match_any_sync, types.i4, types.i8)
+@lower(stubs.match_any_sync, types.i4, types.f4)
+@lower(stubs.match_any_sync, types.i4, types.f8)
+def ptx_match_any_sync(context, builder, sig, args):
+    mask, value = args
+    width = sig.args[1].bitwidth
+    if sig.args[1] in types.real_domain:
+        value = builder.bitcast(value, ir.IntType(width))
+    fname = "llvm.nvvm.match.any.sync.i{}".format(width)
+    lmod = builder.module
+    fnty = ir.FunctionType(ir.IntType(32), (ir.IntType(32), ir.IntType(width)))
+    func = cgutils.get_or_insert_function(lmod, fnty, fname)
+    return builder.call(func, (mask, value))
+
+
+@lower(stubs.match_all_sync, types.i4, types.i4)
+@lower(stubs.match_all_sync, types.i4, types.i8)
+@lower(stubs.match_all_sync, types.i4, types.f4)
+@lower(stubs.match_all_sync, types.i4, types.f8)
+def ptx_match_all_sync(context, builder, sig, args):
+    mask, value = args
+    width = sig.args[1].bitwidth
+    if sig.args[1] in types.real_domain:
+        value = builder.bitcast(value, ir.IntType(width))
+    fname = "llvm.nvvm.match.all.sync.i{}".format(width)
+    lmod = builder.module
+    fnty = ir.FunctionType(
+        ir.LiteralStructType((ir.IntType(32), ir.IntType(1))),
+        (ir.IntType(32), ir.IntType(width)),
+    )
+    func = cgutils.get_or_insert_function(lmod, fnty, fname)
+    return builder.call(func, (mask, value))
+
+
+@lower(stubs.activemask)
+def ptx_activemask(context, builder, sig, args):
+    activemask = ir.InlineAsm(
+        ir.FunctionType(ir.IntType(32), []),
+        "activemask.b32 $0;",
+        "=r",
+        side_effect=True,
+    )
+    return builder.call(activemask, [])
+
+
+@lower(stubs.lanemask_lt)
+def ptx_lanemask_lt(context, builder, sig, args):
+    activemask = ir.InlineAsm(
+        ir.FunctionType(ir.IntType(32), []),
+        "mov.u32 $0, %lanemask_lt;",
+        "=r",
+        side_effect=True,
+    )
+    return builder.call(activemask, [])
+
+
+@lower(stubs.popc, types.Any)
+def ptx_popc(context, builder, sig, args):
+    return builder.ctpop(args[0])
+
+
+@lower(stubs.fma, types.Any, types.Any, types.Any)
+def ptx_fma(context, builder, sig, args):
+    return builder.fma(*args)
+
+
+# See:
+# https://docs.nvidia.com/cuda/libdevice-users-guide/__nv_cbrt.html#__nv_cbrt
+# https://docs.nvidia.com/cuda/libdevice-users-guide/__nv_cbrtf.html#__nv_cbrtf
+
+
+cbrt_funcs = {
+    types.float32: "__nv_cbrtf",
+    types.float64: "__nv_cbrt",
+}
+
+
+@lower(stubs.cbrt, types.float32)
+@lower(stubs.cbrt, types.float64)
+def ptx_cbrt(context, builder, sig, args):
+    ty = sig.return_type
+    fname = cbrt_funcs[ty]
+    fty = context.get_value_type(ty)
+    lmod = builder.module
+    fnty = ir.FunctionType(fty, [fty])
+    fn = cgutils.get_or_insert_function(lmod, fnty, fname)
+    return builder.call(fn, args)
+
+
+@lower(stubs.brev, types.u4)
+def ptx_brev_u4(context, builder, sig, args):
+    # FIXME the llvm.bitreverse.i32 intrinsic isn't supported by nvcc
+    # return builder.bitreverse(args[0])
+
+    fn = cgutils.get_or_insert_function(
+        builder.module,
+        ir.FunctionType(ir.IntType(32), (ir.IntType(32),)),
+        "__nv_brev",
+    )
+    return builder.call(fn, args)
+
+
+@lower(stubs.brev, types.u8)
+def ptx_brev_u8(context, builder, sig, args):
+    # FIXME the llvm.bitreverse.i64 intrinsic isn't supported by nvcc
+    # return builder.bitreverse(args[0])
+
+    fn = cgutils.get_or_insert_function(
+        builder.module,
+        ir.FunctionType(ir.IntType(64), (ir.IntType(64),)),
+        "__nv_brevll",
+    )
+    return builder.call(fn, args)
+
+
+@lower(stubs.clz, types.Any)
+def ptx_clz(context, builder, sig, args):
+    return builder.ctlz(args[0], context.get_constant(types.boolean, 0))
+
+
+@lower(stubs.ffs, types.i4)
+@lower(stubs.ffs, types.u4)
+def ptx_ffs_32(context, builder, sig, args):
+    fn = cgutils.get_or_insert_function(
+        builder.module,
+        ir.FunctionType(ir.IntType(32), (ir.IntType(32),)),
+        "__nv_ffs",
+    )
+    return builder.call(fn, args)
+
+
+@lower(stubs.ffs, types.i8)
+@lower(stubs.ffs, types.u8)
+def ptx_ffs_64(context, builder, sig, args):
+    fn = cgutils.get_or_insert_function(
+        builder.module,
+        ir.FunctionType(ir.IntType(32), (ir.IntType(64),)),
+        "__nv_ffsll",
+    )
+    return builder.call(fn, args)
+
+
+@lower(stubs.selp, types.Any, types.Any, types.Any)
+def ptx_selp(context, builder, sig, args):
+    test, a, b = args
+    return builder.select(test, a, b)
+
+
+@lower(max, types.f4, types.f4)
+def ptx_max_f4(context, builder, sig, args):
+    fn = cgutils.get_or_insert_function(
+        builder.module,
+        ir.FunctionType(ir.FloatType(), (ir.FloatType(), ir.FloatType())),
+        "__nv_fmaxf",
+    )
+    return builder.call(fn, args)
+
+
+@lower(max, types.f8, types.f4)
+@lower(max, types.f4, types.f8)
+@lower(max, types.f8, types.f8)
+def ptx_max_f8(context, builder, sig, args):
+    fn = cgutils.get_or_insert_function(
+        builder.module,
+        ir.FunctionType(ir.DoubleType(), (ir.DoubleType(), ir.DoubleType())),
+        "__nv_fmax",
+    )
+
+    return builder.call(
+        fn,
+        [
+            context.cast(builder, args[0], sig.args[0], types.double),
+            context.cast(builder, args[1], sig.args[1], types.double),
+        ],
+    )
+
+
+@lower(min, types.f4, types.f4)
+def ptx_min_f4(context, builder, sig, args):
+    fn = cgutils.get_or_insert_function(
+        builder.module,
+        ir.FunctionType(ir.FloatType(), (ir.FloatType(), ir.FloatType())),
+        "__nv_fminf",
+    )
+    return builder.call(fn, args)
+
+
+@lower(min, types.f8, types.f4)
+@lower(min, types.f4, types.f8)
+@lower(min, types.f8, types.f8)
+def ptx_min_f8(context, builder, sig, args):
+    fn = cgutils.get_or_insert_function(
+        builder.module,
+        ir.FunctionType(ir.DoubleType(), (ir.DoubleType(), ir.DoubleType())),
+        "__nv_fmin",
+    )
+
+    return builder.call(
+        fn,
+        [
+            context.cast(builder, args[0], sig.args[0], types.double),
+            context.cast(builder, args[1], sig.args[1], types.double),
+        ],
+    )
+
+
+@lower(round, types.f4)
+@lower(round, types.f8)
+def ptx_round(context, builder, sig, args):
+    fn = cgutils.get_or_insert_function(
+        builder.module,
+        ir.FunctionType(ir.IntType(64), (ir.DoubleType(),)),
+        "__nv_llrint",
+    )
+    return builder.call(
+        fn,
+        [
+            context.cast(builder, args[0], sig.args[0], types.double),
+        ],
+    )
+
+
+# This rounding implementation follows the algorithm used in the "fallback
+# version" of double_round in CPython.
+# https://github.com/python/cpython/blob/a755410e054e1e2390de5830befc08fe80706c66/Objects/floatobject.c#L964-L1007
+
+
+@lower(round, types.f4, types.Integer)
+@lower(round, types.f8, types.Integer)
+def round_to_impl(context, builder, sig, args):
+    def round_ndigits(x, ndigits):
+        if math.isinf(x) or math.isnan(x):
+            return x
+
+        if ndigits >= 0:
+            if ndigits > 22:
+                # pow1 and pow2 are each safe from overflow, but
+                # pow1*pow2 ~= pow(10.0, ndigits) might overflow.
+                pow1 = 10.0 ** (ndigits - 22)
+                pow2 = 1e22
+            else:
+                pow1 = 10.0**ndigits
+                pow2 = 1.0
+            y = (x * pow1) * pow2
+            if math.isinf(y):
+                return x
+
+        else:
+            pow1 = 10.0 ** (-ndigits)
+            y = x / pow1
+
+        z = round(y)
+        if math.fabs(y - z) == 0.5:
+            # halfway between two integers; use round-half-even
+            z = 2.0 * round(y / 2.0)
+
+        if ndigits >= 0:
+            z = (z / pow2) / pow1
+        else:
+            z *= pow1
+
+        return z
+
+    return context.compile_internal(
+        builder,
+        round_ndigits,
+        sig,
+        args,
+    )
+
+
+def gen_deg_rad(const):
+    def impl(context, builder, sig, args):
+        (argty,) = sig.args
+        factor = context.get_constant(argty, const)
+        return builder.fmul(factor, args[0])
+
+    return impl
+
+
+_deg2rad = math.pi / 180.0
+_rad2deg = 180.0 / math.pi
+lower(math.radians, types.f4)(gen_deg_rad(_deg2rad))
+lower(math.radians, types.f8)(gen_deg_rad(_deg2rad))
+lower(math.degrees, types.f4)(gen_deg_rad(_rad2deg))
+lower(math.degrees, types.f8)(gen_deg_rad(_rad2deg))
+
+
+def _atomic_dispatcher(dispatch_fn):
+    def imp(context, builder, sig, args):
+        # The common argument handling code
+        aryty, indty, valty = sig.args
+        ary, inds, val = args
+        dtype = aryty.dtype
+
+        indty, indices = normalize_indices(
+            context, builder, indty, inds, aryty, valty
+        )
+
+        lary = context.make_array(aryty)(context, builder, ary)
+        ptr = cgutils.get_item_pointer(
+            context, builder, aryty, lary, indices, wraparound=True
+        )
+        # dispatcher to implementation base on dtype
+        return dispatch_fn(context, builder, dtype, ptr, val)
+
+    return imp
+
+
+@lower(stubs.atomic.add, types.Array, types.intp, types.Any)
+@lower(stubs.atomic.add, types.Array, types.UniTuple, types.Any)
+@lower(stubs.atomic.add, types.Array, types.Tuple, types.Any)
+@_atomic_dispatcher
+def ptx_atomic_add_tuple(context, builder, dtype, ptr, val):
+    if dtype == types.float32:
+        lmod = builder.module
+        return builder.call(
+            nvvmutils.declare_atomic_add_float32(lmod), (ptr, val)
+        )
+    elif dtype == types.float64:
+        lmod = builder.module
+        return builder.call(
+            nvvmutils.declare_atomic_add_float64(lmod), (ptr, val)
+        )
+    else:
+        return builder.atomic_rmw("add", ptr, val, "monotonic")
+
+
+@lower(stubs.atomic.sub, types.Array, types.intp, types.Any)
+@lower(stubs.atomic.sub, types.Array, types.UniTuple, types.Any)
+@lower(stubs.atomic.sub, types.Array, types.Tuple, types.Any)
+@_atomic_dispatcher
+def ptx_atomic_sub(context, builder, dtype, ptr, val):
+    if dtype == types.float32:
+        lmod = builder.module
+        return builder.call(
+            nvvmutils.declare_atomic_sub_float32(lmod), (ptr, val)
+        )
+    elif dtype == types.float64:
+        lmod = builder.module
+        return builder.call(
+            nvvmutils.declare_atomic_sub_float64(lmod), (ptr, val)
+        )
+    else:
+        return builder.atomic_rmw("sub", ptr, val, "monotonic")
+
+
+@lower(stubs.atomic.inc, types.Array, types.intp, types.Any)
+@lower(stubs.atomic.inc, types.Array, types.UniTuple, types.Any)
+@lower(stubs.atomic.inc, types.Array, types.Tuple, types.Any)
+@_atomic_dispatcher
+def ptx_atomic_inc(context, builder, dtype, ptr, val):
+    if dtype in cuda.cudadecl.unsigned_int_numba_types:
+        bw = dtype.bitwidth
+        lmod = builder.module
+        fn = getattr(nvvmutils, f"declare_atomic_inc_int{bw}")
+        return builder.call(fn(lmod), (ptr, val))
+    else:
+        raise TypeError(f"Unimplemented atomic inc with {dtype} array")
+
+
+@lower(stubs.atomic.dec, types.Array, types.intp, types.Any)
+@lower(stubs.atomic.dec, types.Array, types.UniTuple, types.Any)
+@lower(stubs.atomic.dec, types.Array, types.Tuple, types.Any)
+@_atomic_dispatcher
+def ptx_atomic_dec(context, builder, dtype, ptr, val):
+    if dtype in cuda.cudadecl.unsigned_int_numba_types:
+        bw = dtype.bitwidth
+        lmod = builder.module
+        fn = getattr(nvvmutils, f"declare_atomic_dec_int{bw}")
+        return builder.call(fn(lmod), (ptr, val))
+    else:
+        raise TypeError(f"Unimplemented atomic dec with {dtype} array")
+
+
+def ptx_atomic_bitwise(stub, op):
+    @_atomic_dispatcher
+    def impl_ptx_atomic(context, builder, dtype, ptr, val):
+        if dtype in (cuda.cudadecl.integer_numba_types):
+            return builder.atomic_rmw(op, ptr, val, "monotonic")
+        else:
+            raise TypeError(f"Unimplemented atomic {op} with {dtype} array")
+
+    for ty in (types.intp, types.UniTuple, types.Tuple):
+        lower(stub, types.Array, ty, types.Any)(impl_ptx_atomic)
+
+
+ptx_atomic_bitwise(stubs.atomic.and_, "and")
+ptx_atomic_bitwise(stubs.atomic.or_, "or")
+ptx_atomic_bitwise(stubs.atomic.xor, "xor")
+
+
+@lower(stubs.atomic.exch, types.Array, types.intp, types.Any)
+@lower(stubs.atomic.exch, types.Array, types.UniTuple, types.Any)
+@lower(stubs.atomic.exch, types.Array, types.Tuple, types.Any)
+@_atomic_dispatcher
+def ptx_atomic_exch(context, builder, dtype, ptr, val):
+    if dtype in (cuda.cudadecl.integer_numba_types):
+        return builder.atomic_rmw("xchg", ptr, val, "monotonic")
+    else:
+        raise TypeError(f"Unimplemented atomic exch with {dtype} array")
+
+
+@lower(stubs.atomic.max, types.Array, types.intp, types.Any)
+@lower(stubs.atomic.max, types.Array, types.Tuple, types.Any)
+@lower(stubs.atomic.max, types.Array, types.UniTuple, types.Any)
+@_atomic_dispatcher
+def ptx_atomic_max(context, builder, dtype, ptr, val):
+    lmod = builder.module
+    if dtype == types.float64:
+        return builder.call(
+            nvvmutils.declare_atomic_max_float64(lmod), (ptr, val)
+        )
+    elif dtype == types.float32:
+        return builder.call(
+            nvvmutils.declare_atomic_max_float32(lmod), (ptr, val)
+        )
+    elif dtype in (types.int32, types.int64):
+        return builder.atomic_rmw("max", ptr, val, ordering="monotonic")
+    elif dtype in (types.uint32, types.uint64):
+        return builder.atomic_rmw("umax", ptr, val, ordering="monotonic")
+    else:
+        raise TypeError("Unimplemented atomic max with %s array" % dtype)
+
+
+@lower(stubs.atomic.min, types.Array, types.intp, types.Any)
+@lower(stubs.atomic.min, types.Array, types.Tuple, types.Any)
+@lower(stubs.atomic.min, types.Array, types.UniTuple, types.Any)
+@_atomic_dispatcher
+def ptx_atomic_min(context, builder, dtype, ptr, val):
+    lmod = builder.module
+    if dtype == types.float64:
+        return builder.call(
+            nvvmutils.declare_atomic_min_float64(lmod), (ptr, val)
+        )
+    elif dtype == types.float32:
+        return builder.call(
+            nvvmutils.declare_atomic_min_float32(lmod), (ptr, val)
+        )
+    elif dtype in (types.int32, types.int64):
+        return builder.atomic_rmw("min", ptr, val, ordering="monotonic")
+    elif dtype in (types.uint32, types.uint64):
+        return builder.atomic_rmw("umin", ptr, val, ordering="monotonic")
+    else:
+        raise TypeError("Unimplemented atomic min with %s array" % dtype)
+
+
+@lower(stubs.atomic.nanmax, types.Array, types.intp, types.Any)
+@lower(stubs.atomic.nanmax, types.Array, types.Tuple, types.Any)
+@lower(stubs.atomic.nanmax, types.Array, types.UniTuple, types.Any)
+@_atomic_dispatcher
+def ptx_atomic_nanmax(context, builder, dtype, ptr, val):
+    lmod = builder.module
+    if dtype == types.float64:
+        return builder.call(
+            nvvmutils.declare_atomic_nanmax_float64(lmod), (ptr, val)
+        )
+    elif dtype == types.float32:
+        return builder.call(
+            nvvmutils.declare_atomic_nanmax_float32(lmod), (ptr, val)
+        )
+    elif dtype in (types.int32, types.int64):
+        return builder.atomic_rmw("max", ptr, val, ordering="monotonic")
+    elif dtype in (types.uint32, types.uint64):
+        return builder.atomic_rmw("umax", ptr, val, ordering="monotonic")
+    else:
+        raise TypeError("Unimplemented atomic max with %s array" % dtype)
+
+
+@lower(stubs.atomic.nanmin, types.Array, types.intp, types.Any)
+@lower(stubs.atomic.nanmin, types.Array, types.Tuple, types.Any)
+@lower(stubs.atomic.nanmin, types.Array, types.UniTuple, types.Any)
+@_atomic_dispatcher
+def ptx_atomic_nanmin(context, builder, dtype, ptr, val):
+    lmod = builder.module
+    if dtype == types.float64:
+        return builder.call(
+            nvvmutils.declare_atomic_nanmin_float64(lmod), (ptr, val)
+        )
+    elif dtype == types.float32:
+        return builder.call(
+            nvvmutils.declare_atomic_nanmin_float32(lmod), (ptr, val)
+        )
+    elif dtype in (types.int32, types.int64):
+        return builder.atomic_rmw("min", ptr, val, ordering="monotonic")
+    elif dtype in (types.uint32, types.uint64):
+        return builder.atomic_rmw("umin", ptr, val, ordering="monotonic")
+    else:
+        raise TypeError("Unimplemented atomic min with %s array" % dtype)
+
+
+@lower(stubs.atomic.compare_and_swap, types.Array, types.Any, types.Any)
+def ptx_atomic_compare_and_swap(context, builder, sig, args):
+    sig = sig.return_type(sig.args[0], types.intp, sig.args[1], sig.args[2])
+    args = (args[0], context.get_constant(types.intp, 0), args[1], args[2])
+    return ptx_atomic_cas(context, builder, sig, args)
+
+
+@lower(stubs.atomic.cas, types.Array, types.intp, types.Any, types.Any)
+@lower(stubs.atomic.cas, types.Array, types.Tuple, types.Any, types.Any)
+@lower(stubs.atomic.cas, types.Array, types.UniTuple, types.Any, types.Any)
+def ptx_atomic_cas(context, builder, sig, args):
+    aryty, indty, oldty, valty = sig.args
+    ary, inds, old, val = args
+
+    indty, indices = normalize_indices(
+        context, builder, indty, inds, aryty, valty
+    )
+
+    lary = context.make_array(aryty)(context, builder, ary)
+    ptr = cgutils.get_item_pointer(
+        context, builder, aryty, lary, indices, wraparound=True
+    )
+
+    if aryty.dtype in (cuda.cudadecl.integer_numba_types):
+        lmod = builder.module
+        bitwidth = aryty.dtype.bitwidth
+        return nvvmutils.atomic_cmpxchg(builder, lmod, bitwidth, ptr, old, val)
+    else:
+        raise TypeError("Unimplemented atomic cas with %s array" % aryty.dtype)
+
+
+# -----------------------------------------------------------------------------
+
+
+@lower(breakpoint)
+def ptx_brkpt(context, builder, sig, args):
+    brkpt = ir.InlineAsm(
+        ir.FunctionType(ir.VoidType(), []), "brkpt;", "", side_effect=True
+    )
+    builder.call(brkpt, ())
+
+
+@lower(stubs.nanosleep, types.uint32)
+def ptx_nanosleep(context, builder, sig, args):
+    nanosleep = ir.InlineAsm(
+        ir.FunctionType(ir.VoidType(), [ir.IntType(32)]),
+        "nanosleep.u32 $0;",
+        "r",
+        side_effect=True,
+    )
+    ns = args[0]
+    builder.call(nanosleep, [ns])
+
+
+# -----------------------------------------------------------------------------
+
+
+def _generic_array(
+    context,
+    builder,
+    shape,
+    dtype,
+    symbol_name,
+    addrspace,
+    can_dynsized=False,
+    alignment=None,
+):
+    elemcount = reduce(operator.mul, shape, 1)
+
+    # Check for valid shape for this type of allocation.
+    # Only 1d arrays can be dynamic.
+    dynamic_smem = elemcount <= 0 and can_dynsized and len(shape) == 1
+    if elemcount <= 0 and not dynamic_smem:
+        raise ValueError("array length <= 0")
+
+    # Check that we support the requested dtype
+    number_domain = types.number_domain
+    struct_model_types = (StructModel,)
+    if cuda.HAS_NUMBA:
+        number_domain |= core_types.number_domain
+        struct_model_types = (StructModel, CoreStructModel)
+
+    data_model = context.data_model_manager[dtype]
+    other_supported_type = (
+        isinstance(dtype, (types.Record, types.Boolean))
+        or isinstance(data_model, struct_model_types)
+        or dtype == types.float16
+    )
+    if dtype not in number_domain and not other_supported_type:
+        raise TypeError("unsupported type: %s" % dtype)
+
+    lldtype = context.get_data_type(dtype)
+    laryty = ir.ArrayType(lldtype, elemcount)
+
+    if addrspace == nvvm.ADDRSPACE_LOCAL:
+        # Special case local address space allocation to use alloca
+        # NVVM is smart enough to only use local memory if no register is
+        # available
+        dataptr = cgutils.alloca_once(builder, laryty, name=symbol_name)
+
+        # If the caller has specified a custom alignment, just set the align
+        # attribute on the alloca IR directly.  We don't do any additional
+        # hand-holding here like checking the underlying data type's alignment
+        # or rounding up to the next power of 2--those checks will have already
+        # been done by the time we see the alignment value.
+        if alignment is not None:
+            dataptr.align = alignment
+    else:
+        lmod = builder.module
+
+        # Create global variable in the requested address space
+        gvmem = cgutils.add_global_variable(
+            lmod, laryty, symbol_name, addrspace
+        )
+
+        # If the caller hasn't specified a custom alignment, obtain the
+        # underlying dtype alignment from the ABI and then round it up to
+        # a power of two.  Otherwise, just use the caller's alignment.
+        #
+        # N.B. The caller *could* provide a valid-but-smaller-than-natural
+        #      alignment here; we'll assume the caller knows what they're
+        #      doing and let that through without error.
+
+        if alignment is None:
+            abi_alignment = context.get_abi_alignment(lldtype)
+            # Alignment is required to be a power of 2 for shared memory.
+            # If it is not a power of 2 (e.g. for a Record array) then round
+            # up accordingly.
+            actual_alignment = 1 << (abi_alignment - 1).bit_length()
+        else:
+            actual_alignment = alignment
+
+        gvmem.align = actual_alignment
+
+        if dynamic_smem:
+            gvmem.linkage = "external"
+        else:
+            ## Comment out the following line to workaround a NVVM bug
+            ## which generates a invalid symbol name when the linkage
+            ## is internal and in some situation.
+            ## See _get_unique_smem_id()
+            # gvmem.linkage = lc.LINKAGE_INTERNAL
+
+            gvmem.initializer = ir.Constant(laryty, ir.Undefined)
+
+        # Convert to generic address-space
+        dataptr = builder.addrspacecast(
+            gvmem, ir.PointerType(ir.IntType(8)), "generic"
+        )
+
+    targetdata = ll.create_target_data(nvvm.NVVM().data_layout)
+    lldtype = context.get_data_type(dtype)
+    itemsize = lldtype.get_abi_size(targetdata)
+
+    # Compute strides
+    laststride = itemsize
+    rstrides = []
+    for i, lastsize in enumerate(reversed(shape)):
+        rstrides.append(laststride)
+        laststride *= lastsize
+    strides = [s for s in reversed(rstrides)]
+    kstrides = [context.get_constant(types.intp, s) for s in strides]
+
+    # Compute shape
+    if dynamic_smem:
+        # Compute the shape based on the dynamic shared memory configuration.
+        # Unfortunately NVVM does not provide an intrinsic for the
+        # %dynamic_smem_size register, so we must read it using inline
+        # assembly.
+        get_dynshared_size = ir.InlineAsm(
+            ir.FunctionType(ir.IntType(32), []),
+            "mov.u32 $0, %dynamic_smem_size;",
+            "=r",
+            side_effect=True,
+        )
+        dynsmem_size = builder.zext(
+            builder.call(get_dynshared_size, []), ir.IntType(64)
+        )
+        # Only 1-D dynamic shared memory is supported so the following is a
+        # sufficient construction of the shape
+        kitemsize = context.get_constant(types.intp, itemsize)
+        kshape = [builder.udiv(dynsmem_size, kitemsize)]
+    else:
+        kshape = [context.get_constant(types.intp, s) for s in shape]
+
+    # Create array object
+    ndim = len(shape)
+    aryty = types.Array(dtype=dtype, ndim=ndim, layout="C")
+    ary = context.make_array(aryty)(context, builder)
+
+    context.populate_array(
+        ary,
+        data=builder.bitcast(dataptr, ary.data.type),
+        shape=kshape,
+        strides=kstrides,
+        itemsize=context.get_constant(types.intp, itemsize),
+        meminfo=None,
+    )
+    return ary._getvalue()
+
+
+@lower_constant(CUDADispatcher)
+def cuda_dispatcher_const(context, builder, ty, pyval):
+    return context.get_dummy_value()
+
+
+# NumPy
+
+register_ufuncs(ufunc_db.get_ufuncs(), lower)