numba-cuda 0.0.1__py3-none-any.whl → 0.0.13__py3-none-any.whl

numba_cuda/numba/cuda/cudaimpl.py

@@ -0,0 +1,1055 @@
+from functools import reduce
+import operator
+import math
+
+from llvmlite import ir
+import llvmlite.binding as ll
+
+from numba.core.imputils import Registry, lower_cast
+from numba.core.typing.npydecl import parse_dtype
+from numba.core.datamodel import models
+from numba.core import types, cgutils
+from numba.np import ufunc_db
+from numba.np.npyimpl import register_ufuncs
+from .cudadrv import nvvm
+from numba import cuda
+from numba.cuda import nvvmutils, stubs, errors
+from numba.cuda.types import dim3, CUDADispatcher
+
+registry = Registry()
+lower = registry.lower
+lower_attr = registry.lower_getattr
+lower_constant = registry.lower_constant
+
+
+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)
+
+
+@lower(cuda.shared.array, types.IntegerLiteral, types.Any)
+def cuda_shared_array_integer(context, builder, sig, args):
+    length = sig.args[0].literal_value
+    dtype = parse_dtype(sig.args[1])
+    return _generic_array(context, builder, shape=(length,), dtype=dtype,
+                          symbol_name=_get_unique_smem_id('_cudapy_smem'),
+                          addrspace=nvvm.ADDRSPACE_SHARED,
+                          can_dynsized=True)
+
+
+@lower(cuda.shared.array, types.Tuple, types.Any)
+@lower(cuda.shared.array, types.UniTuple, types.Any)
+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])
+    return _generic_array(context, builder, shape=shape, dtype=dtype,
+                          symbol_name=_get_unique_smem_id('_cudapy_smem'),
+                          addrspace=nvvm.ADDRSPACE_SHARED,
+                          can_dynsized=True)
+
+
+@lower(cuda.local.array, types.IntegerLiteral, types.Any)
+def cuda_local_array_integer(context, builder, sig, args):
+    length = sig.args[0].literal_value
+    dtype = parse_dtype(sig.args[1])
+    return _generic_array(context, builder, shape=(length,), dtype=dtype,
+                          symbol_name='_cudapy_lmem',
+                          addrspace=nvvm.ADDRSPACE_LOCAL,
+                          can_dynsized=False)
+
+
+@lower(cuda.local.array, types.Tuple, types.Any)
+@lower(cuda.local.array, types.UniTuple, types.Any)
+def ptx_lmem_alloc_array(context, builder, sig, args):
+    shape = [ s.literal_value for s in sig.args[0] ]
+    dtype = parse_dtype(sig.args[1])
+    return _generic_array(context, builder, shape=shape, dtype=dtype,
+                          symbol_name='_cudapy_lmem',
+                          addrspace=nvvm.ADDRSPACE_LOCAL,
+                          can_dynsized=False)
+
+
+@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.shfl_sync_intrinsic, types.i4, types.i4, types.i4, types.i4,
+       types.i4)
+@lower(stubs.shfl_sync_intrinsic, types.i4, types.i4, types.i8, types.i4,
+       types.i4)
+@lower(stubs.shfl_sync_intrinsic, types.i4, types.i4, types.f4, types.i4,
+       types.i4)
+@lower(stubs.shfl_sync_intrinsic, types.i4, types.i4, types.f8, types.i4,
+       types.i4)
+def ptx_shfl_sync_i32(context, builder, sig, args):
+    """
+    The NVVM intrinsic for shfl only supports i32, but the cuda intrinsic
+    function supports both 32 and 64 bit ints and floats, so for feature parity,
+    i64, f32, and f64 are implemented. Floats by way of bitcasting the float to
+    an int, then shuffling, then bitcasting back. And 64-bit values by packing
+    them into 2 32bit values, shuffling thoose, and then packing back together.
+    """
+    mask, mode, value, index, clamp = args
+    value_type = sig.args[2]
+    if value_type in types.real_domain:
+        value = builder.bitcast(value, ir.IntType(value_type.bitwidth))
+    fname = 'llvm.nvvm.shfl.sync.i32'
+    lmod = builder.module
+    fnty = ir.FunctionType(
+        ir.LiteralStructType((ir.IntType(32), ir.IntType(1))),
+                            (ir.IntType(32), ir.IntType(32), ir.IntType(32),
+                             ir.IntType(32), ir.IntType(32))
+    )
+    func = cgutils.get_or_insert_function(lmod, fnty, fname)
+    if value_type.bitwidth == 32:
+        ret = builder.call(func, (mask, mode, value, index, clamp))
+        if value_type == types.float32:
+            rv = builder.extract_value(ret, 0)
+            pred = builder.extract_value(ret, 1)
+            fv = builder.bitcast(rv, ir.FloatType())
+            ret = cgutils.make_anonymous_struct(builder, (fv, pred))
+    else:
+        value1 = builder.trunc(value, ir.IntType(32))
+        value_lshr = builder.lshr(value, context.get_constant(types.i8, 32))
+        value2 = builder.trunc(value_lshr, ir.IntType(32))
+        ret1 = builder.call(func, (mask, mode, value1, index, clamp))
+        ret2 = builder.call(func, (mask, mode, value2, index, clamp))
+        rv1 = builder.extract_value(ret1, 0)
+        rv2 = builder.extract_value(ret2, 0)
+        pred = builder.extract_value(ret1, 1)
+        rv1_64 = builder.zext(rv1, ir.IntType(64))
+        rv2_64 = builder.zext(rv2, ir.IntType(64))
+        rv_shl = builder.shl(rv2_64, context.get_constant(types.i8, 32))
+        rv = builder.or_(rv_shl, rv1_64)
+        if value_type == types.float64:
+            rv = builder.bitcast(rv, ir.DoubleType())
+        ret = cgutils.make_anonymous_struct(builder, (rv, pred))
+    return ret
+
+
+@lower(stubs.vote_sync_intrinsic, types.i4, types.i4, types.boolean)
+def ptx_vote_sync(context, builder, sig, args):
+    fname = 'llvm.nvvm.vote.sync'
+    lmod = builder.module
+    fnty = ir.FunctionType(ir.LiteralStructType((ir.IntType(32),
+                                                 ir.IntType(1))),
+                           (ir.IntType(32), ir.IntType(32), ir.IntType(1)))
+    func = cgutils.get_or_insert_function(lmod, fnty, fname)
+    return builder.call(func, args)
+
+
+@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)
+
+
+def float16_float_ty_constraint(bitwidth):
+    typemap = {32: ('f32', 'f'), 64: ('f64', 'd')}
+
+    try:
+        return typemap[bitwidth]
+    except KeyError:
+        msg = f"Conversion between float16 and float{bitwidth} unsupported"
+        raise errors.CudaLoweringError(msg)
+
+
+@lower_cast(types.float16, types.Float)
+def float16_to_float_cast(context, builder, fromty, toty, val):
+    if fromty.bitwidth == toty.bitwidth:
+        return val
+
+    ty, constraint = float16_float_ty_constraint(toty.bitwidth)
+
+    fnty = ir.FunctionType(context.get_value_type(toty), [ir.IntType(16)])
+    asm = ir.InlineAsm(fnty, f"cvt.{ty}.f16 $0, $1;", f"={constraint},h")
+    return builder.call(asm, [val])
+
+
+@lower_cast(types.Float, types.float16)
+def float_to_float16_cast(context, builder, fromty, toty, val):
+    if fromty.bitwidth == toty.bitwidth:
+        return val
+
+    ty, constraint = float16_float_ty_constraint(fromty.bitwidth)
+
+    fnty = ir.FunctionType(ir.IntType(16), [context.get_value_type(fromty)])
+    asm = ir.InlineAsm(fnty, f"cvt.rn.f16.{ty} $0, $1;", f"=h,{constraint}")
+    return builder.call(asm, [val])
+
+
+def float16_int_constraint(bitwidth):
+    typemap = { 8: 'c', 16: 'h', 32: 'r', 64: 'l' }
+
+    try:
+        return typemap[bitwidth]
+    except KeyError:
+        msg = f"Conversion between float16 and int{bitwidth} unsupported"
+        raise errors.CudaLoweringError(msg)
+
+
+@lower_cast(types.float16, types.Integer)
+def float16_to_integer_cast(context, builder, fromty, toty, val):
+    bitwidth = toty.bitwidth
+    constraint = float16_int_constraint(bitwidth)
+    signedness = 's' if toty.signed else 'u'
+
+    fnty = ir.FunctionType(context.get_value_type(toty), [ir.IntType(16)])
+    asm = ir.InlineAsm(fnty,
+                       f"cvt.rni.{signedness}{bitwidth}.f16 $0, $1;",
+                       f"={constraint},h")
+    return builder.call(asm, [val])
+
+
+@lower_cast(types.Integer, types.float16)
+@lower_cast(types.IntegerLiteral, types.float16)
+def integer_to_float16_cast(context, builder, fromty, toty, val):
+    bitwidth = fromty.bitwidth
+    constraint = float16_int_constraint(bitwidth)
+    signedness = 's' if fromty.signed else 'u'
+
+    fnty = ir.FunctionType(ir.IntType(16),
+                           [context.get_value_type(fromty)])
+    asm = ir.InlineAsm(fnty,
+                       f"cvt.rn.f16.{signedness}{bitwidth} $0, $1;",
+                       f"=h,{constraint}")
+    return builder.call(asm, [val])
+
+
+def lower_fp16_binary(fn, op):
+    @lower(fn, types.float16, types.float16)
+    def ptx_fp16_binary(context, builder, sig, args):
+        fnty = ir.FunctionType(ir.IntType(16),
+                               [ir.IntType(16), ir.IntType(16)])
+        asm = ir.InlineAsm(fnty, f'{op}.f16 $0,$1,$2;', '=h,h,h')
+        return builder.call(asm, args)
+
+
+lower_fp16_binary(stubs.fp16.hadd, 'add')
+lower_fp16_binary(operator.add, 'add')
+lower_fp16_binary(operator.iadd, 'add')
+lower_fp16_binary(stubs.fp16.hsub, 'sub')
+lower_fp16_binary(operator.sub, 'sub')
+lower_fp16_binary(operator.isub, 'sub')
+lower_fp16_binary(stubs.fp16.hmul, 'mul')
+lower_fp16_binary(operator.mul, 'mul')
+lower_fp16_binary(operator.imul, 'mul')
+
+
+@lower(stubs.fp16.hneg, types.float16)
+def ptx_fp16_hneg(context, builder, sig, args):
+    fnty = ir.FunctionType(ir.IntType(16), [ir.IntType(16)])
+    asm = ir.InlineAsm(fnty, 'neg.f16 $0, $1;', '=h,h')
+    return builder.call(asm, args)
+
+
+@lower(operator.neg, types.float16)
+def operator_hneg(context, builder, sig, args):
+    return ptx_fp16_hneg(context, builder, sig, args)
+
+
+@lower(stubs.fp16.habs, types.float16)
+def ptx_fp16_habs(context, builder, sig, args):
+    fnty = ir.FunctionType(ir.IntType(16), [ir.IntType(16)])
+    asm = ir.InlineAsm(fnty, 'abs.f16 $0, $1;', '=h,h')
+    return builder.call(asm, args)
+
+
+@lower(abs, types.float16)
+def operator_habs(context, builder, sig, args):
+    return ptx_fp16_habs(context, builder, sig, args)
+
+
+@lower(stubs.fp16.hfma, types.float16, types.float16, types.float16)
+def ptx_hfma(context, builder, sig, args):
+    argtys = [ir.IntType(16), ir.IntType(16), ir.IntType(16)]
+    fnty = ir.FunctionType(ir.IntType(16), argtys)
+    asm = ir.InlineAsm(fnty, "fma.rn.f16 $0,$1,$2,$3;", "=h,h,h,h")
+    return builder.call(asm, args)
+
+
+@lower(operator.truediv, types.float16, types.float16)
+@lower(operator.itruediv, types.float16, types.float16)
+def fp16_div_impl(context, builder, sig, args):
+    def fp16_div(x, y):
+        return cuda.fp16.hdiv(x, y)
+
+    return context.compile_internal(builder, fp16_div, sig, args)
+
+
+_fp16_cmp = """{{
+          .reg .pred __$$f16_cmp_tmp;
+          setp.{op}.f16 __$$f16_cmp_tmp, $1, $2;
+          selp.u16 $0, 1, 0, __$$f16_cmp_tmp;
+        }}"""
+
+
+def _gen_fp16_cmp(op):
+    def ptx_fp16_comparison(context, builder, sig, args):
+        fnty = ir.FunctionType(ir.IntType(16), [ir.IntType(16), ir.IntType(16)])
+        asm = ir.InlineAsm(fnty, _fp16_cmp.format(op=op), '=h,h,h')
+        result = builder.call(asm, args)
+
+        zero = context.get_constant(types.int16, 0)
+        int_result = builder.bitcast(result, ir.IntType(16))
+        return builder.icmp_unsigned("!=", int_result, zero)
+    return ptx_fp16_comparison
+
+
+lower(stubs.fp16.heq, types.float16, types.float16)(_gen_fp16_cmp('eq'))
+lower(operator.eq, types.float16, types.float16)(_gen_fp16_cmp('eq'))
+lower(stubs.fp16.hne, types.float16, types.float16)(_gen_fp16_cmp('ne'))
+lower(operator.ne, types.float16, types.float16)(_gen_fp16_cmp('ne'))
+lower(stubs.fp16.hge, types.float16, types.float16)(_gen_fp16_cmp('ge'))
+lower(operator.ge, types.float16, types.float16)(_gen_fp16_cmp('ge'))
+lower(stubs.fp16.hgt, types.float16, types.float16)(_gen_fp16_cmp('gt'))
+lower(operator.gt, types.float16, types.float16)(_gen_fp16_cmp('gt'))
+lower(stubs.fp16.hle, types.float16, types.float16)(_gen_fp16_cmp('le'))
+lower(operator.le, types.float16, types.float16)(_gen_fp16_cmp('le'))
+lower(stubs.fp16.hlt, types.float16, types.float16)(_gen_fp16_cmp('lt'))
+lower(operator.lt, types.float16, types.float16)(_gen_fp16_cmp('lt'))
+
+
+def lower_fp16_minmax(fn, fname, op):
+    @lower(fn, types.float16, types.float16)
+    def ptx_fp16_minmax(context, builder, sig, args):
+        choice = _gen_fp16_cmp(op)(context, builder, sig, args)
+        return builder.select(choice, args[0], args[1])
+
+
+lower_fp16_minmax(stubs.fp16.hmax, 'max', 'gt')
+lower_fp16_minmax(stubs.fp16.hmin, 'min', 'lt')
+
+# 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.
+_rad2deg = 180. / 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 _normalize_indices(context, builder, indty, inds, aryty, valty):
+    """
+    Convert integer indices into tuple of intp
+    """
+    if indty in types.integer_domain:
+        indty = types.UniTuple(dtype=indty, count=1)
+        indices = [inds]
+    else:
+        indices = cgutils.unpack_tuple(builder, inds, count=len(indty))
+    indices = [context.cast(builder, i, t, types.intp)
+               for t, i in zip(indty, indices)]
+
+    dtype = aryty.dtype
+    if dtype != valty:
+        raise TypeError("expect %s but got %s" % (dtype, valty))
+
+    if aryty.ndim != len(indty):
+        raise TypeError("indexing %d-D array with %d-D index" %
+                        (aryty.ndim, len(indty)))
+
+    return indty, indices
+
+
+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(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):
+    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
+    data_model = context.data_model_manager[dtype]
+    other_supported_type = (
+        isinstance(dtype, (types.Record, types.Boolean))
+        or isinstance(data_model, models.StructModel)
+        or dtype == types.float16
+    )
+    if dtype not in types.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)
+    else:
+        lmod = builder.module
+
+        # Create global variable in the requested address space
+        gvmem = cgutils.add_global_variable(lmod, laryty, symbol_name,
+                                            addrspace)
+        # Specify alignment to avoid misalignment bug
+        align = context.get_abi_sizeof(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.
+        gvmem.align = 1 << (align - 1 ).bit_length()
+
+        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)