numba-cuda 0.8.1__py3-none-any.whl → 0.10.0__py3-none-any.whl

numba_cuda/numba/cuda/initialize.py

@@ -4,9 +4,11 @@ def initialize_all():
 
     from numba.cuda.decorators import jit
     from numba.cuda.dispatcher import CUDADispatcher
-    from numba.core.target_extension import (target_registry,
-                                             dispatcher_registry,
-                                             jit_registry)
+    from numba.core.target_extension import (
+        target_registry,
+        dispatcher_registry,
+        jit_registry,
+    )
 
     cuda_target = target_registry["cuda"]
     jit_registry[cuda_target] = jit

numba_cuda/numba/cuda/intrinsic_wrapper.py

@@ -36,42 +36,3 @@ def ballot_sync(mask, predicate):
     and are within the given mask.
     """
     return numba.cuda.vote_sync_intrinsic(mask, 3, predicate)[0]
-
-
-@jit(device=True)
-def shfl_sync(mask, value, src_lane):
-    """
-    Shuffles value across the masked warp and returns the value
-    from src_lane. If this is outside the warp, then the
-    given value is returned.
-    """
-    return numba.cuda.shfl_sync_intrinsic(mask, 0, value, src_lane, 0x1f)[0]
-
-
-@jit(device=True)
-def shfl_up_sync(mask, value, delta):
-    """
-    Shuffles value across the masked warp and returns the value
-    from (laneid - delta). If this is outside the warp, then the
-    given value is returned.
-    """
-    return numba.cuda.shfl_sync_intrinsic(mask, 1, value, delta, 0)[0]
-
-
-@jit(device=True)
-def shfl_down_sync(mask, value, delta):
-    """
-    Shuffles value across the masked warp and returns the value
-    from (laneid + delta). If this is outside the warp, then the
-    given value is returned.
-    """
-    return numba.cuda.shfl_sync_intrinsic(mask, 2, value, delta, 0x1f)[0]
-
-
-@jit(device=True)
-def shfl_xor_sync(mask, value, lane_mask):
-    """
-    Shuffles value across the masked warp and returns the value
-    from (laneid ^ lane_mask).
-    """
-    return numba.cuda.shfl_sync_intrinsic(mask, 3, value, lane_mask, 0x1f)[0]

numba_cuda/numba/cuda/intrinsics.py

@@ -2,16 +2,17 @@ from llvmlite import ir
 
 from numba import cuda, types
 from numba.core import cgutils
-from numba.core.errors import RequireLiteralValue
+from numba.core.errors import RequireLiteralValue, TypingError
 from numba.core.typing import signature
 from numba.core.extending import overload_attribute, overload_method
 from numba.cuda import nvvmutils
 from numba.cuda.extending import intrinsic
 
 
-#-------------------------------------------------------------------------------
+# -------------------------------------------------------------------------------
 # Grid functions
 
+
 def _type_grid_function(ndim):
     val = ndim.literal_value
     if val == 1:
@@ -19,14 +20,14 @@ def _type_grid_function(ndim):
     elif val in (2, 3):
         restype = types.UniTuple(types.int64, val)
     else:
-        raise ValueError('argument can only be 1, 2, 3')
+        raise ValueError("argument can only be 1, 2, 3")
 
     return signature(restype, types.int32)
 
 
 @intrinsic
 def grid(typingctx, ndim):
-    '''grid(ndim)
+    """grid(ndim)
 
     Return the absolute position of the current thread in the entire grid of
     blocks.  *ndim* should correspond to the number of dimensions declared when
@@ -39,7 +40,7 @@ def grid(typingctx, ndim):
 
     and is similar for the other two indices, but using the ``y`` and ``z``
     attributes.
-    '''
+    """
 
     if not isinstance(ndim, types.IntegerLiteral):
         raise RequireLiteralValue(ndim)
@@ -59,7 +60,7 @@ def grid(typingctx, ndim):
 
 @intrinsic
 def gridsize(typingctx, ndim):
-    '''gridsize(ndim)
+    """gridsize(ndim)
 
     Return the absolute size (or shape) in threads of the entire grid of
     blocks. *ndim* should correspond to the number of dimensions declared when
@@ -72,7 +73,7 @@ def gridsize(typingctx, ndim):
 
     and is similar for the other two indices, but using the ``y`` and ``z``
     attributes.
-    '''
+    """
 
     if not isinstance(ndim, types.IntegerLiteral):
         raise RequireLiteralValue(ndim)
@@ -87,17 +88,17 @@ def gridsize(typingctx, ndim):
 
     def codegen(context, builder, sig, args):
         restype = sig.return_type
-        nx = _nthreads_for_dim(builder, 'x')
+        nx = _nthreads_for_dim(builder, "x")
 
         if restype == types.int64:
             return nx
         elif isinstance(restype, types.UniTuple):
-            ny = _nthreads_for_dim(builder, 'y')
+            ny = _nthreads_for_dim(builder, "y")
 
             if restype.count == 2:
                 return cgutils.pack_array(builder, (nx, ny))
             elif restype.count == 3:
-                nz = _nthreads_for_dim(builder, 'z')
+                nz = _nthreads_for_dim(builder, "z")
                 return cgutils.pack_array(builder, (nx, ny, nz))
 
     return sig, codegen
@@ -108,37 +109,40 @@ def _warpsize(typingctx):
     sig = signature(types.int32)
 
     def codegen(context, builder, sig, args):
-        return nvvmutils.call_sreg(builder, 'warpsize')
+        return nvvmutils.call_sreg(builder, "warpsize")
 
     return sig, codegen
 
 
-@overload_attribute(types.Module(cuda), 'warpsize', target='cuda')
+@overload_attribute(types.Module(cuda), "warpsize", target="cuda")
 def cuda_warpsize(mod):
-    '''
+    """
     The size of a warp. All architectures implemented to date have a warp size
     of 32.
-    '''
+    """
+
     def get(mod):
         return _warpsize()
+
     return get
 
 
-#-------------------------------------------------------------------------------
+# -------------------------------------------------------------------------------
 # syncthreads
 
+
 @intrinsic
 def syncthreads(typingctx):
-    '''
+    """
     Synchronize all threads in the same thread block.  This function implements
     the same pattern as barriers in traditional multi-threaded programming: this
     function waits until all threads in the block call it, at which point it
     returns control to all its callers.
-    '''
+    """
     sig = signature(types.none)
 
     def codegen(context, builder, sig, args):
-        fname = 'llvm.nvvm.barrier0'
+        fname = "llvm.nvvm.barrier0"
         lmod = builder.module
         fnty = ir.FunctionType(ir.VoidType(), ())
         sync = cgutils.get_or_insert_function(lmod, fnty, fname)
@@ -164,40 +168,211 @@ def _syncthreads_predicate(typingctx, predicate, fname):
 
 @intrinsic
 def syncthreads_count(typingctx, predicate):
-    '''
+    """
     syncthreads_count(predicate)
 
     An extension to numba.cuda.syncthreads where the return value is a count
     of the threads where predicate is true.
-    '''
-    fname = 'llvm.nvvm.barrier0.popc'
+    """
+    fname = "llvm.nvvm.barrier0.popc"
     return _syncthreads_predicate(typingctx, predicate, fname)
 
 
 @intrinsic
 def syncthreads_and(typingctx, predicate):
-    '''
+    """
     syncthreads_and(predicate)
 
     An extension to numba.cuda.syncthreads where 1 is returned if predicate is
     true for all threads or 0 otherwise.
-    '''
-    fname = 'llvm.nvvm.barrier0.and'
+    """
+    fname = "llvm.nvvm.barrier0.and"
     return _syncthreads_predicate(typingctx, predicate, fname)
 
 
 @intrinsic
 def syncthreads_or(typingctx, predicate):
-    '''
+    """
     syncthreads_or(predicate)
 
     An extension to numba.cuda.syncthreads where 1 is returned if predicate is
     true for any thread or 0 otherwise.
-    '''
-    fname = 'llvm.nvvm.barrier0.or'
+    """
+    fname = "llvm.nvvm.barrier0.or"
     return _syncthreads_predicate(typingctx, predicate, fname)
 
 
-@overload_method(types.Integer, 'bit_count', target='cuda')
+@overload_method(types.Integer, "bit_count", target="cuda")
 def integer_bit_count(i):
     return lambda i: cuda.popc(i)
+
+
+# -------------------------------------------------------------------------------
+# Warp shuffle functions
+#
+# References:
+#
+# - https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#warp-shuffle-functions
+# - https://docs.nvidia.com/cuda/nvvm-ir-spec/index.html#data-movement
+#
+# Notes:
+#
+# - The public CUDA C/C++ and Numba Python APIs for these intrinsics use
+#   different names for parameters to the NVVM IR specification. So that we
+#   can correlate the implementation with the documentation, the @intrinsic
+#   API functions map the public API arguments to the NVVM intrinsic
+#   arguments.
+# - The NVVM IR specification requires some of the parameters (e.g. mode) to be
+#   constants. It's therefore essential that we pass in some values to the
+#   shfl_sync_intrinsic function (e.g. the mode and c values).
+# - Normally parameters for intrinsic functions in Numba would be given the
+#   same name as used in the API, and would contain a type. However, because we
+#   have to pass in some values and some times (and there is divergence between
+#   the names in the intrinsic documentation and the public APIs) we instead
+#   follow the convention of naming shfl_sync_intrinsic parameters with a
+#   suffix of _type or _value depending on whether they contain a type or a
+#   value.
+
+
+@intrinsic
+def shfl_sync(typingctx, mask, value, src_lane):
+    """
+    Shuffles ``value`` across the masked warp and returns the value from
+    ``src_lane``. If this is outside the warp, then the given value is
+    returned.
+    """
+    membermask_type = mask
+    mode_value = 0
+    a_type = value
+    b_type = src_lane
+    c_value = 0x1F
+    return shfl_sync_intrinsic(
+        typingctx, membermask_type, mode_value, a_type, b_type, c_value
+    )
+
+
+@intrinsic
+def shfl_up_sync(typingctx, mask, value, delta):
+    """
+    Shuffles ``value`` across the masked warp and returns the value from
+    ``(laneid - delta)``. If this is outside the warp, then the given value is
+    returned.
+    """
+    membermask_type = mask
+    mode_value = 1
+    a_type = value
+    b_type = delta
+    c_value = 0
+    return shfl_sync_intrinsic(
+        typingctx, membermask_type, mode_value, a_type, b_type, c_value
+    )
+
+
+@intrinsic
+def shfl_down_sync(typingctx, mask, value, delta):
+    """
+    Shuffles ``value`` across the masked warp and returns the value from
+    ``(laneid + delta)``. If this is outside the warp, then the given value is
+    returned.
+    """
+    membermask_type = mask
+    mode_value = 2
+    a_type = value
+    b_type = delta
+    c_value = 0x1F
+    return shfl_sync_intrinsic(
+        typingctx, membermask_type, mode_value, a_type, b_type, c_value
+    )
+
+
+@intrinsic
+def shfl_xor_sync(typingctx, mask, value, lane_mask):
+    """
+    Shuffles ``value`` across the masked warp and returns the value from
+    ``(laneid ^ lane_mask)``.
+    """
+    membermask_type = mask
+    mode_value = 3
+    a_type = value
+    b_type = lane_mask
+    c_value = 0x1F
+    return shfl_sync_intrinsic(
+        typingctx, membermask_type, mode_value, a_type, b_type, c_value
+    )
+
+
+def shfl_sync_intrinsic(
+    typingctx,
+    membermask_type,
+    mode_value,
+    a_type,
+    b_type,
+    c_value,
+):
+    if a_type not in (types.i4, types.i8, types.f4, types.f8):
+        raise TypingError(
+            "shfl_sync only supports 32- and 64-bit ints and floats"
+        )
+
+    def codegen(context, builder, sig, args):
+        """
+        The NVVM shfl_sync intrinsic only supports i32, but the CUDA C/C++
+        intrinsic supports both 32- and 64-bit ints and floats, so for feature
+        parity, i32, i64, f32, and f64 are implemented. Floats by way of
+        bitcasting the float to an int, then shuffling, then bitcasting
+        back."""
+        membermask, a, b = args
+
+        # Types
+        a_type = sig.args[1]
+        return_type = context.get_value_type(sig.return_type)
+        i32 = ir.IntType(32)
+        i64 = ir.IntType(64)
+
+        if a_type in types.real_domain:
+            a = builder.bitcast(a, ir.IntType(a_type.bitwidth))
+
+        # NVVM intrinsic definition
+        arg_types = (i32, i32, i32, i32, i32)
+        shfl_return_type = ir.LiteralStructType((i32, ir.IntType(1)))
+        fnty = ir.FunctionType(shfl_return_type, arg_types)
+
+        fname = "llvm.nvvm.shfl.sync.i32"
+        shfl_sync = cgutils.get_or_insert_function(builder.module, fnty, fname)
+
+        # Intrinsic arguments
+        mode = ir.Constant(i32, mode_value)
+        c = ir.Constant(i32, c_value)
+        membermask = builder.trunc(membermask, i32)
+        b = builder.trunc(b, i32)
+
+        if a_type.bitwidth == 32:
+            a = builder.trunc(a, i32)
+            ret = builder.call(shfl_sync, (membermask, mode, a, b, c))
+            d = builder.extract_value(ret, 0)
+        else:
+            # Handle 64-bit values by shuffling as two 32-bit values and
+            # packing the result into 64 bits.
+
+            # Extract high and low parts
+            lo = builder.trunc(a, i32)
+            a_lshr = builder.lshr(a, ir.Constant(i64, 32))
+            hi = builder.trunc(a_lshr, i32)
+
+            # Shuffle individual parts
+            ret_lo = builder.call(shfl_sync, (membermask, mode, lo, b, c))
+            ret_hi = builder.call(shfl_sync, (membermask, mode, hi, b, c))
+
+            # Combine individual result parts into a 64-bit result
+            d_lo = builder.extract_value(ret_lo, 0)
+            d_hi = builder.extract_value(ret_hi, 0)
+            d_lo_64 = builder.zext(d_lo, i64)
+            d_hi_64 = builder.zext(d_hi, i64)
+            d_shl = builder.shl(d_hi_64, ir.Constant(i64, 32))
+            d = builder.or_(d_shl, d_lo_64)
+
+        return builder.bitcast(d, return_type)
+
+    sig = signature(a_type, membermask_type, a_type, b_type)
+
+    return sig, codegen

numba_cuda/numba/cuda/kernels/reduction.py

@@ -13,7 +13,7 @@ def _gpu_reduce_factory(fn, nbtype):
     from numba import cuda
 
     reduce_op = cuda.jit(device=True)(fn)
-    inner_sm_size = _WARPSIZE + 1   # plus one to avoid SM collision
+    inner_sm_size = _WARPSIZE + 1  # plus one to avoid SM collision
     max_blocksize = _NUMWARPS * _WARPSIZE
 
     @cuda.jit(device=True)
@@ -86,8 +86,9 @@ def _gpu_reduce_factory(fn, nbtype):
         # warning: this is assuming 4 warps.
         # assert numwarps == 4
         if tid < 2:
-            sm_partials[tid, 0] = reduce_op(sm_partials[tid, 0],
-                                            sm_partials[tid + 2, 0])
+            sm_partials[tid, 0] = reduce_op(
+                sm_partials[tid, 0], sm_partials[tid + 2, 0]
+            )
             cuda.syncwarp()
         if tid == 0:
             partials[blkid] = reduce_op(sm_partials[0, 0], sm_partials[1, 0])
@@ -148,8 +149,9 @@ def _gpu_reduce_factory(fn, nbtype):
         """
         tid = cuda.threadIdx.x
 
-        sm_partials = cuda.shared.array((_NUMWARPS, inner_sm_size),
-                                        dtype=nbtype)
+        sm_partials = cuda.shared.array(
+            (_NUMWARPS, inner_sm_size), dtype=nbtype
+        )
         if cuda.blockDim.x == max_blocksize:
             device_reduce_full_block(arr, partials, sm_partials)
         else:
@@ -238,17 +240,15 @@ class Reduce(object):
 
         if size_full:
             # kernel for the fully populated threadblocks
-            kernel[full_blockct, blocksize, stream](arr[:size_full],
-                                                    partials[:full_blockct],
-                                                    init,
-                                                    True)
+            kernel[full_blockct, blocksize, stream](
+                arr[:size_full], partials[:full_blockct], init, True
+            )
 
         if size_partial:
             # kernel for partially populated threadblocks
-            kernel[1, size_partial, stream](arr[size_full:],
-                                            partials[full_blockct:],
-                                            init,
-                                            not full_blockct)
+            kernel[1, size_partial, stream](
+                arr[size_full:], partials[full_blockct:], init, not full_blockct
+            )
 
         if partials.size > 1:
             # finish up

numba_cuda/numba/cuda/kernels/transpose.py

@@ -18,16 +18,14 @@ def transpose(a, b=None):
     """
 
     # prefer `a`'s stream if
-    stream = getattr(a, 'stream', 0)
+    stream = getattr(a, "stream", 0)
 
     if not b:
         cols, rows = a.shape
         strides = a.dtype.itemsize * cols, a.dtype.itemsize
         b = cuda.cudadrv.devicearray.DeviceNDArray(
-            (rows, cols),
-            strides,
-            dtype=a.dtype,
-            stream=stream)
+            (rows, cols), strides, dtype=a.dtype, stream=stream
+        )
 
     dt = nps.from_dtype(a.dtype)
 
@@ -40,7 +38,6 @@ def transpose(a, b=None):
 
     @cuda.jit
     def kernel(input, output):
-
         tile = cuda.shared.array(shape=tile_shape, dtype=dt)
 
         tx = cuda.threadIdx.x