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

numba_cuda/numba/cuda/stubs.py

@@ -1,6 +1,7 @@
 """
 This scripts specifies all PTX special objects.
 """
+
 import numpy as np
 from collections import defaultdict
 import functools
@@ -9,12 +10,13 @@ from inspect import Signature, Parameter
 
 
 class Stub(object):
-    '''
+    """
     A stub object to represent special objects that are meaningless
     outside the context of a CUDA kernel
-    '''
-    _description_ = '<ptx special value>'
-    __slots__ = () # don't allocate __dict__
+    """
+
+    _description_ = "<ptx special value>"
+    __slots__ = ()  # don't allocate __dict__
 
     def __new__(cls):
         raise NotImplementedError("%s is not instantiable" % cls)
@@ -24,23 +26,26 @@ class Stub(object):
 
 
 def stub_function(fn):
-    '''
+    """
     A stub function to represent special functions that are meaningless
     outside the context of a CUDA kernel
-    '''
+    """
+
     @functools.wraps(fn)
     def wrapped(*args, **kwargs):
         raise NotImplementedError("%s cannot be called from host code" % fn)
+
     return wrapped
 
 
-#-------------------------------------------------------------------------------
+# -------------------------------------------------------------------------------
 # Thread and grid indices and dimensions
 
 
 class Dim3(Stub):
-    '''A triple, (x, y, z)'''
-    _description_ = '<Dim3>'
+    """A triple, (x, y, z)"""
+
+    _description_ = "<Dim3>"
 
     @property
     def x(self):
@@ -56,68 +61,76 @@ class Dim3(Stub):
 
 
 class threadIdx(Dim3):
-    '''
+    """
     The thread indices in the current thread block. Each index is an integer
     spanning the range from 0 inclusive to the corresponding value of the
     attribute in :attr:`numba.cuda.blockDim` exclusive.
-    '''
-    _description_ = '<threadIdx.{x,y,z}>'
+    """
+
+    _description_ = "<threadIdx.{x,y,z}>"
 
 
 class blockIdx(Dim3):
-    '''
+    """
     The block indices in the grid of thread blocks. Each index is an integer
     spanning the range from 0 inclusive to the corresponding value of the
     attribute in :attr:`numba.cuda.gridDim` exclusive.
-    '''
-    _description_ = '<blockIdx.{x,y,z}>'
+    """
+
+    _description_ = "<blockIdx.{x,y,z}>"
 
 
 class blockDim(Dim3):
-    '''
+    """
     The shape of a block of threads, as declared when instantiating the kernel.
     This value is the same for all threads in a given kernel launch, even if
     they belong to different blocks (i.e. each block is "full").
-    '''
-    _description_ = '<blockDim.{x,y,z}>'
+    """
+
+    _description_ = "<blockDim.{x,y,z}>"
 
 
 class gridDim(Dim3):
-    '''
+    """
     The shape of the grid of blocks. This value is the same for all threads in
     a given kernel launch.
-    '''
-    _description_ = '<gridDim.{x,y,z}>'
+    """
+
+    _description_ = "<gridDim.{x,y,z}>"
 
 
 class warpsize(Stub):
-    '''
+    """
     The size of a warp. All architectures implemented to date have a warp size
     of 32.
-    '''
-    _description_ = '<warpsize>'
+    """
+
+    _description_ = "<warpsize>"
 
 
 class laneid(Stub):
-    '''
+    """
     This thread's lane within a warp. Ranges from 0 to
     :attr:`numba.cuda.warpsize` - 1.
-    '''
-    _description_ = '<laneid>'
+    """
+
+    _description_ = "<laneid>"
 
 
-#-------------------------------------------------------------------------------
+# -------------------------------------------------------------------------------
 # Array creation
 
+
 class shared(Stub):
-    '''
+    """
     Shared memory namespace
-    '''
-    _description_ = '<shared>'
+    """
+
+    _description_ = "<shared>"
 
     @stub_function
     def array(shape, dtype):
-        '''
+        """
         Allocate a shared array of the given *shape* and *type*. *shape* is
         either an integer or a tuple of integers representing the array's
         dimensions.  *type* is a :ref:`Numba type <numba-types>` of the
@@ -125,83 +138,89 @@ class shared(Stub):
 
         The returned array-like object can be read and written to like any
         normal device array (e.g. through indexing).
-        '''
+        """
 
 
 class local(Stub):
-    '''
+    """
     Local memory namespace
-    '''
-    _description_ = '<local>'
+    """
+
+    _description_ = "<local>"
 
     @stub_function
     def array(shape, dtype):
-        '''
+        """
         Allocate a local array of the given *shape* and *type*. The array is
         private to the current thread, and resides in global memory. An
         array-like object is returned which can be read and written to like any
         standard array (e.g.  through indexing).
-        '''
+        """
 
 
 class const(Stub):
-    '''
+    """
     Constant memory namespace
-    '''
+    """
 
     @stub_function
     def array_like(ndarray):
-        '''
+        """
         Create a const array from *ndarry*. The resulting const array will have
         the same shape, type, and values as *ndarray*.
-        '''
+        """
 
 
 # -------------------------------------------------------------------------------
 # warp level operations
 
+
 class syncwarp(Stub):
-    '''
+    """
     syncwarp(mask=0xFFFFFFFF)
 
     Synchronizes a masked subset of threads in a warp.
-    '''
-    _description_ = '<warp_sync()>'
+    """
+
+    _description_ = "<warp_sync()>"
 
 
 class shfl_sync_intrinsic(Stub):
-    '''
+    """
     shfl_sync_intrinsic(mask, mode, value, mode_offset, clamp)
 
     Nvvm intrinsic for shuffling data across a warp
     docs.nvidia.com/cuda/nvvm-ir-spec/index.html#nvvm-intrin-warp-level-datamove
-    '''
-    _description_ = '<shfl_sync()>'
+    """
+
+    _description_ = "<shfl_sync()>"
 
 
 class vote_sync_intrinsic(Stub):
-    '''
+    """
     vote_sync_intrinsic(mask, mode, predictate)
 
     Nvvm intrinsic for performing a reduce and broadcast across a warp
     docs.nvidia.com/cuda/nvvm-ir-spec/index.html#nvvm-intrin-warp-level-vote
-    '''
-    _description_ = '<vote_sync()>'
+    """
+
+    _description_ = "<vote_sync()>"
 
 
 class match_any_sync(Stub):
-    '''
+    """
     match_any_sync(mask, value)
 
     Nvvm intrinsic for performing a compare and broadcast across a warp.
     Returns a mask of threads that have same value as the given value from
     within the masked warp.
-    '''
-    _description_ = '<match_any_sync()>'
+    """
+
+    _description_ = "<match_any_sync()>"
 
 
 class match_all_sync(Stub):
-    '''
+    """
     match_all_sync(mask, value)
 
     Nvvm intrinsic for performing a compare and broadcast across a warp.
@@ -209,12 +228,13 @@ class match_all_sync(Stub):
     same value as the given value from within the masked warp, if they
     all have the same value, otherwise it is 0. Pred is a boolean of whether
     or not all threads in the mask warp have the same warp.
-    '''
-    _description_ = '<match_all_sync()>'
+    """
+
+    _description_ = "<match_all_sync()>"
 
 
 class activemask(Stub):
-    '''
+    """
     activemask()
 
     Returns a 32-bit integer mask of all currently active threads in the
@@ -222,47 +242,54 @@ class activemask(Stub):
     activemask() is called. Inactive threads are represented by 0 bits in the
     returned mask. Threads which have exited the kernel are always marked as
     inactive.
-    '''
-    _description_ = '<activemask()>'
+    """
+
+    _description_ = "<activemask()>"
 
 
 class lanemask_lt(Stub):
-    '''
+    """
     lanemask_lt()
 
     Returns a 32-bit integer mask of all lanes (including inactive ones) with
     ID less than the current lane.
-    '''
-    _description_ = '<lanemask_lt()>'
+    """
+
+    _description_ = "<lanemask_lt()>"
 
 
 # -------------------------------------------------------------------------------
 # memory fences
 
+
 class threadfence_block(Stub):
-    '''
+    """
     A memory fence at thread block level
-    '''
-    _description_ = '<threadfence_block()>'
+    """
+
+    _description_ = "<threadfence_block()>"
 
 
 class threadfence_system(Stub):
-    '''
+    """
     A memory fence at system level: across devices
-    '''
-    _description_ = '<threadfence_system()>'
+    """
+
+    _description_ = "<threadfence_system()>"
 
 
 class threadfence(Stub):
-    '''
+    """
     A memory fence at device level
-    '''
-    _description_ = '<threadfence()>'
+    """
+
+    _description_ = "<threadfence()>"
 
 
-#-------------------------------------------------------------------------------
+# -------------------------------------------------------------------------------
 # bit manipulation
 
+
 class popc(Stub):
     """
     popc(x)
@@ -297,9 +324,10 @@ class ffs(Stub):
     """
 
 
-#-------------------------------------------------------------------------------
+# -------------------------------------------------------------------------------
 # comparison and selection instructions
 
+
 class selp(Stub):
     """
     selp(a, b, c)
@@ -309,9 +337,10 @@ class selp(Stub):
     """
 
 
-#-------------------------------------------------------------------------------
+# -------------------------------------------------------------------------------
 # single / double precision arithmetic
 
+
 class fma(Stub):
     """
     fma(a, b, c)
@@ -321,20 +350,21 @@ class fma(Stub):
 
 
 class cbrt(Stub):
-    """"
+    """ "
     cbrt(a)
 
     Perform the cube root operation.
     """
 
 
-#-------------------------------------------------------------------------------
+# -------------------------------------------------------------------------------
 # atomic
 
+
 class atomic(Stub):
-    """Namespace for atomic operations
-    """
-    _description_ = '<atomic>'
+    """Namespace for atomic operations"""
+
+    _description_ = "<atomic>"
 
     class add(Stub):
         """add(ary, idx, val)
@@ -401,8 +431,7 @@ class atomic(Stub):
 
         Performs::
 
-           ary[idx] = (val if (ary[idx] == 0) or
-                       (ary[idx] > val) else ary[idx] - 1)
+           ary[idx] = val if (ary[idx] == 0) or (ary[idx] > val) else ary[idx] - 1
 
         Supported on uint32, and uint64 operands only.
 
@@ -497,26 +526,29 @@ class atomic(Stub):
         """
 
 
-#-------------------------------------------------------------------------------
+# -------------------------------------------------------------------------------
 # timers
 
+
 class nanosleep(Stub):
-    '''
+    """
     nanosleep(ns)
 
     Suspends the thread for a sleep duration approximately close to the delay
     `ns`, specified in nanoseconds.
-    '''
-    _description_ = '<nansleep()>'
+    """
+
+    _description_ = "<nansleep()>"
 
-#-------------------------------------------------------------------------------
+
+# -------------------------------------------------------------------------------
 # Floating point 16
 
 
 class fp16(Stub):
-    """Namespace for fp16 operations
-    """
-    _description_ = '<fp16>'
+    """Namespace for fp16 operations"""
+
+    _description_ = "<fp16>"
 
     class hadd(Stub):
         """hadd(a, b)
@@ -817,9 +849,10 @@ class fp16(Stub):
         """
 
 
-#-------------------------------------------------------------------------------
+# -------------------------------------------------------------------------------
 # vector types
 
+
 def make_vector_type_stubs():
     """Make user facing objects for vector types"""
     vector_type_stubs = []
@@ -833,7 +866,7 @@ def make_vector_type_stubs():
         "uint32",
         "uint64",
         "float32",
-        "float64"
+        "float64",
     )
     vector_type_element_counts = (1, 2, 3, 4)
     vector_type_attribute_names = ("x", "y", "z", "w")
@@ -845,21 +878,25 @@ def make_vector_type_stubs():
         attr_names = vector_type_attribute_names[:nelem]
 
         vector_type_stub = type(
-            type_name, (Stub,),
+            type_name,
+            (Stub,),
             {
                 **{attr: lambda self: None for attr in attr_names},
                 **{
                     "_description_": f"<{type_name}>",
-                    "__signature__": Signature(parameters=[
-                        Parameter(
-                            name=attr_name, kind=Parameter.POSITIONAL_ONLY
-                        ) for attr_name in attr_names[:nelem]
-                    ]),
+                    "__signature__": Signature(
+                        parameters=[
+                            Parameter(
+                                name=attr_name, kind=Parameter.POSITIONAL_ONLY
+                            )
+                            for attr_name in attr_names[:nelem]
+                        ]
+                    ),
                     "__doc__": f"A stub for {type_name} to be used in "
-                    "CUDA kernels."
+                    "CUDA kernels.",
                 },
-                **{"aliases": []}
-            }
+                **{"aliases": []},
+            },
         )
         vector_type_stubs.append(vector_type_stub)
     return vector_type_stubs
@@ -884,7 +921,7 @@ def map_vector_type_stubs_to_alias(vector_type_stubs):
         "ulong": f"uint{np.dtype(np.uint).itemsize * 8}",
         "ulonglong": f"uint{np.dtype(np.ulonglong).itemsize * 8}",
         "float": f"float{np.dtype(np.single).itemsize * 8}",
-        "double": f"float{np.dtype(np.double).itemsize * 8}"
+        "double": f"float{np.dtype(np.double).itemsize * 8}",
     }
 
     base_type_to_vector_type = defaultdict(list)