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

numba_cuda/numba/cuda/random.py

@@ -1,7 +1,16 @@
 import math
 
-from numba import (config, cuda, float32, float64, uint32, int64, uint64,
-                   from_dtype, jit)
+from numba import (
+    config,
+    cuda,
+    float32,
+    float64,
+    uint32,
+    int64,
+    uint64,
+    from_dtype,
+    jit,
+)
 
 import numpy as np
 
@@ -29,8 +38,9 @@ import numpy as np
 # using the CPU @jit decorator everywhere to create functions that work as
 # both CPU and CUDA device functions.
 
-xoroshiro128p_dtype = np.dtype([('s0', np.uint64), ('s1', np.uint64)],
-                               align=True)
+xoroshiro128p_dtype = np.dtype(
+    [("s0", np.uint64), ("s1", np.uint64)], align=True
+)
 xoroshiro128p_type = from_dtype(xoroshiro128p_dtype)
 
 # When cudasim is enabled, Fake CUDA arrays are passed to some of the
@@ -45,7 +55,7 @@ _nopython = not config.ENABLE_CUDASIM
 
 @jit(forceobj=_forceobj, looplift=_looplift, nopython=_nopython)
 def init_xoroshiro128p_state(states, index, seed):
-    '''Use SplitMix64 to generate an xoroshiro128p state from 64-bit seed.
+    """Use SplitMix64 to generate an xoroshiro128p state from 64-bit seed.
 
     This ensures that manually set small seeds don't result in a predictable
     initial sequence from the random number generator.
@@ -56,7 +66,7 @@ def init_xoroshiro128p_state(states, index, seed):
     :param index: offset in states to update
     :type seed: int64
     :param seed: seed value to use when initializing state
-    '''
+    """
     index = int64(index)
     seed = uint64(seed)
 
@@ -65,13 +75,13 @@ def init_xoroshiro128p_state(states, index, seed):
     z = (z ^ (z >> uint32(27))) * uint64(0x94D049BB133111EB)
     z = z ^ (z >> uint32(31))
 
-    states[index]['s0'] = z
-    states[index]['s1'] = z
+    states[index]["s0"] = z
+    states[index]["s1"] = z
 
 
 @jit(forceobj=_forceobj, looplift=_looplift, nopython=_nopython)
 def rotl(x, k):
-    '''Left rotate x by k bits.'''
+    """Left rotate x by k bits."""
     x = uint64(x)
     k = uint32(k)
     return (x << k) | (x >> uint32(64 - k))
@@ -79,38 +89,38 @@ def rotl(x, k):
 
 @jit(forceobj=_forceobj, looplift=_looplift, nopython=_nopython)
 def xoroshiro128p_next(states, index):
-    '''Return the next random uint64 and advance the RNG in states[index].
+    """Return the next random uint64 and advance the RNG in states[index].
 
     :type states: 1D array, dtype=xoroshiro128p_dtype
     :param states: array of RNG states
     :type index: int64
     :param index: offset in states to update
     :rtype: uint64
-    '''
+    """
     index = int64(index)
-    s0 = states[index]['s0']
-    s1 = states[index]['s1']
+    s0 = states[index]["s0"]
+    s1 = states[index]["s1"]
     result = s0 + s1
 
     s1 ^= s0
-    states[index]['s0'] = uint64(rotl(s0, uint32(55))) ^ s1 ^ (s1 << uint32(14))
-    states[index]['s1'] = uint64(rotl(s1, uint32(36)))
+    states[index]["s0"] = uint64(rotl(s0, uint32(55))) ^ s1 ^ (s1 << uint32(14))
+    states[index]["s1"] = uint64(rotl(s1, uint32(36)))
 
     return result
 
 
 @jit(forceobj=_forceobj, looplift=_looplift, nopython=_nopython)
 def xoroshiro128p_jump(states, index):
-    '''Advance the RNG in ``states[index]`` by 2**64 steps.
+    """Advance the RNG in ``states[index]`` by 2**64 steps.
 
     :type states: 1D array, dtype=xoroshiro128p_dtype
     :param states: array of RNG states
     :type index: int64
     :param index: offset in states to update
-    '''
+    """
     index = int64(index)
 
-    jump = (uint64(0xbeac0467eba5facb), uint64(0xd86b048b86aa9922))
+    jump = (uint64(0xBEAC0467EBA5FACB), uint64(0xD86B048B86AA9922))
 
     s0 = uint64(0)
     s1 = uint64(0)
@@ -118,52 +128,52 @@ def xoroshiro128p_jump(states, index):
     for i in range(2):
         for b in range(64):
             if jump[i] & (uint64(1) << uint32(b)):
-                s0 ^= states[index]['s0']
-                s1 ^= states[index]['s1']
+                s0 ^= states[index]["s0"]
+                s1 ^= states[index]["s1"]
             xoroshiro128p_next(states, index)
 
-    states[index]['s0'] = s0
-    states[index]['s1'] = s1
+    states[index]["s0"] = s0
+    states[index]["s1"] = s1
 
 
 @jit(forceobj=_forceobj, looplift=_looplift, nopython=_nopython)
 def uint64_to_unit_float64(x):
-    '''Convert uint64 to float64 value in the range [0.0, 1.0)'''
+    """Convert uint64 to float64 value in the range [0.0, 1.0)"""
     x = uint64(x)
     return (x >> uint32(11)) * (float64(1) / (uint64(1) << uint32(53)))
 
 
 @jit(forceobj=_forceobj, looplift=_looplift, nopython=_nopython)
 def uint64_to_unit_float32(x):
-    '''Convert uint64 to float32 value in the range [0.0, 1.0)'''
+    """Convert uint64 to float32 value in the range [0.0, 1.0)"""
     x = uint64(x)
     return float32(uint64_to_unit_float64(x))
 
 
 @jit(forceobj=_forceobj, looplift=_looplift, nopython=_nopython)
 def xoroshiro128p_uniform_float32(states, index):
-    '''Return a float32 in range [0.0, 1.0) and advance ``states[index]``.
+    """Return a float32 in range [0.0, 1.0) and advance ``states[index]``.
 
     :type states: 1D array, dtype=xoroshiro128p_dtype
     :param states: array of RNG states
     :type index: int64
     :param index: offset in states to update
     :rtype: float32
-    '''
+    """
     index = int64(index)
     return uint64_to_unit_float32(xoroshiro128p_next(states, index))
 
 
 @jit(forceobj=_forceobj, looplift=_looplift, nopython=_nopython)
 def xoroshiro128p_uniform_float64(states, index):
-    '''Return a float64 in range [0.0, 1.0) and advance ``states[index]``.
+    """Return a float64 in range [0.0, 1.0) and advance ``states[index]``.
 
     :type states: 1D array, dtype=xoroshiro128p_dtype
     :param states: array of RNG states
     :type index: int64
     :param index: offset in states to update
     :rtype: float64
-    '''
+    """
     index = int64(index)
     return uint64_to_unit_float64(xoroshiro128p_next(states, index))
 
@@ -174,7 +184,7 @@ TWO_PI_FLOAT64 = np.float64(2 * math.pi)
 
 @jit(forceobj=_forceobj, looplift=_looplift, nopython=_nopython)
 def xoroshiro128p_normal_float32(states, index):
-    '''Return a normally distributed float32 and advance ``states[index]``.
+    """Return a normally distributed float32 and advance ``states[index]``.
 
     The return value is drawn from a Gaussian of mean=0 and sigma=1 using the
     Box-Muller transform.  This advances the RNG sequence by two steps.
@@ -184,7 +194,7 @@ def xoroshiro128p_normal_float32(states, index):
     :type index: int64
     :param index: offset in states to update
     :rtype: float32
-    '''
+    """
     index = int64(index)
 
     u1 = xoroshiro128p_uniform_float32(states, index)
@@ -199,7 +209,7 @@ def xoroshiro128p_normal_float32(states, index):
 
 @jit(forceobj=_forceobj, looplift=_looplift, nopython=_nopython)
 def xoroshiro128p_normal_float64(states, index):
-    '''Return a normally distributed float32 and advance ``states[index]``.
+    """Return a normally distributed float32 and advance ``states[index]``.
 
     The return value is drawn from a Gaussian of mean=0 and sigma=1 using the
     Box-Muller transform.  This advances the RNG sequence by two steps.
@@ -209,7 +219,7 @@ def xoroshiro128p_normal_float64(states, index):
     :type index: int64
     :param index: offset in states to update
     :rtype: float64
-    '''
+    """
     index = int64(index)
 
     u1 = xoroshiro128p_uniform_float32(states, index)
@@ -242,7 +252,7 @@ def init_xoroshiro128p_states_cpu(states, seed, subsequence_start):
 
 
 def init_xoroshiro128p_states(states, seed, subsequence_start=0, stream=0):
-    '''Initialize RNG states on the GPU for parallel generators.
+    """Initialize RNG states on the GPU for parallel generators.
 
     This initializes the RNG states so that each state in the array corresponds
     subsequences in the separated by 2**64 steps from each other in the main
@@ -257,7 +267,7 @@ def init_xoroshiro128p_states(states, seed, subsequence_start=0, stream=0):
     :param states: array of RNG states
     :type seed: uint64
     :param seed: starting seed for list of generators
-    '''
+    """
 
     # Initialization on CPU is much faster than the GPU
     states_cpu = np.empty(shape=states.shape, dtype=xoroshiro128p_dtype)
@@ -267,7 +277,7 @@ def init_xoroshiro128p_states(states, seed, subsequence_start=0, stream=0):
 
 
 def create_xoroshiro128p_states(n, seed, subsequence_start=0, stream=0):
-    '''Returns a new device array initialized for n random number generators.
+    """Returns a new device array initialized for n random number generators.
 
     This initializes the RNG states so that each state in the array corresponds
     subsequences in the separated by 2**64 steps from each other in the main
@@ -286,7 +296,7 @@ def create_xoroshiro128p_states(n, seed, subsequence_start=0, stream=0):
     :param subsequence_start:
     :type stream: CUDA stream
     :param stream: stream to run initialization kernel on
-    '''
+    """
     states = cuda.device_array(n, dtype=xoroshiro128p_dtype, stream=stream)
     init_xoroshiro128p_states(states, seed, subsequence_start, stream)
     return states

numba_cuda/numba/cuda/reshape_funcs.cu

@@ -148,4 +148,4 @@ numba_attempt_nocopy_reshape(npy_intp nd, const npy_intp *dims, const npy_intp *
     }
 
     return 1;
-}
+}

numba_cuda/numba/cuda/runtime/__init__.py

@@ -1 +1 @@
-from numba.cuda.runtime.nrt import rtsys # noqa: F401
+from numba.cuda.runtime.nrt import rtsys  # noqa: F401

numba_cuda/numba/cuda/runtime/memsys.cu

@@ -91,4 +91,4 @@ extern "C" __global__ void NRT_MemSys_print(void)
   } else {
     printf("TheMsys is null.\n");
   }
-}
+}

numba_cuda/numba/cuda/runtime/memsys.cuh

@@ -14,4 +14,4 @@ struct NRT_MemSys {
 /* The Memory System object */
 __device__ NRT_MemSys* TheMSys;
 
-extern "C" __global__ void NRT_MemSys_set(NRT_MemSys *memsys_ptr);
+extern "C" __global__ void NRT_MemSys_set(NRT_MemSys *memsys_ptr);

numba_cuda/numba/cuda/runtime/nrt.cu

@@ -33,7 +33,7 @@ extern "C" __device__ void* NRT_Allocate(size_t size)
 {
   void* ptr = NULL;
   ptr       = malloc(size);
-  if (TheMSys && TheMSys->stats.enabled) { 
+  if (TheMSys && TheMSys->stats.enabled) {
     TheMSys->stats.alloc.fetch_add(1, cuda::memory_order_relaxed); }
   return ptr;
 }
@@ -49,7 +49,7 @@ extern "C" __device__ void NRT_MemInfo_init(NRT_MemInfo* mi,
   mi->dtor_info = dtor_info;
   mi->data      = data;
   mi->size      = size;
- if (TheMSys && TheMSys->stats.enabled) { 
+ if (TheMSys && TheMSys->stats.enabled) {
   TheMSys->stats.mi_alloc.fetch_add(1, cuda::memory_order_relaxed); }
 }
 
@@ -77,7 +77,7 @@ extern "C" __device__ void NRT_dealloc(NRT_MemInfo* mi)
 extern "C" __device__ void NRT_MemInfo_destroy(NRT_MemInfo* mi)
 {
   NRT_dealloc(mi);
-  if (TheMSys && TheMSys->stats.enabled) { 
+  if (TheMSys && TheMSys->stats.enabled) {
     TheMSys->stats.mi_free.fetch_add(1, cuda::memory_order_relaxed); }
 }
 

numba_cuda/numba/cuda/runtime/nrt.py

@@ -5,26 +5,28 @@ import numpy as np
 
 from numba import cuda, config
 from numba.core.runtime.nrt import _nrt_mstats
-from numba.cuda.cudadrv.driver import (Linker, driver, launch_kernel,
-                                       USE_NV_BINDING)
+from numba.cuda.cudadrv.driver import (
+    Linker,
+    driver,
+    launch_kernel,
+    USE_NV_BINDING,
+)
 from numba.cuda.cudadrv import devices
 from numba.cuda.api import get_current_device
 from numba.cuda.utils import _readenv
 
 
 # Check environment variable or config for NRT statistics enablement
-NRT_STATS = (
-    _readenv("NUMBA_CUDA_NRT_STATS", bool, False) or
-    getattr(config, "NUMBA_CUDA_NRT_STATS", False)
+NRT_STATS = _readenv("NUMBA_CUDA_NRT_STATS", bool, False) or getattr(
+    config, "NUMBA_CUDA_NRT_STATS", False
 )
 if not hasattr(config, "NUMBA_CUDA_NRT_STATS"):
     config.CUDA_NRT_STATS = NRT_STATS
 
 
 # Check environment variable or config for NRT enablement
-ENABLE_NRT = (
-    _readenv("NUMBA_CUDA_ENABLE_NRT", bool, False) or
-    getattr(config, "NUMBA_CUDA_ENABLE_NRT", False)
+ENABLE_NRT = _readenv("NUMBA_CUDA_ENABLE_NRT", bool, False) or getattr(
+    config, "NUMBA_CUDA_ENABLE_NRT", False
 )
 if not hasattr(config, "NUMBA_CUDA_ENABLE_NRT"):
     config.CUDA_ENABLE_NRT = ENABLE_NRT
@@ -35,16 +37,19 @@ def _alloc_init_guard(method):
     """
     Ensure NRT memory allocation and initialization before running the method
     """
+
     @wraps(method)
     def wrapper(self, *args, **kwargs):
         self.ensure_allocated()
         self.ensure_initialized()
         return method(self, *args, **kwargs)
+
     return wrapper
 
 
 class _Runtime:
     """Singleton class for Numba CUDA runtime"""
+
     _instance = None
 
     def __new__(cls, *args, **kwargs):
@@ -64,8 +69,7 @@ class _Runtime:
         """
         # Define the path for memsys.cu
         memsys_mod = os.path.join(
-            os.path.dirname(os.path.abspath(__file__)),
-            "memsys.cu"
+            os.path.dirname(os.path.abspath(__file__)), "memsys.cu"
         )
         cc = get_current_device().compute_capability
 
@@ -105,10 +109,12 @@ class _Runtime:
         # Allocate space for NRT_MemSys
         ptr, nbytes = self._memsys_module.get_global_symbol("memsys_size")
         memsys_size = ctypes.c_uint64()
-        driver.cuMemcpyDtoH(ctypes.addressof(memsys_size),
-                            ptr.device_ctypes_pointer, nbytes)
+        driver.cuMemcpyDtoH(
+            ctypes.addressof(memsys_size), ptr.device_ctypes_pointer, nbytes
+        )
         self._memsys = device_array(
-            (memsys_size.value,), dtype="i1", stream=stream)
+            (memsys_size.value,), dtype="i1", stream=stream
+        )
         self.set_memsys_to_module(self._memsys_module, stream=stream)
 
     def _single_thread_launch(self, module, stream, name, params=()):
@@ -121,12 +127,16 @@ class _Runtime:
         func = module.get_function(name)
         launch_kernel(
             func.handle,
-            1, 1, 1,
-            1, 1, 1,
+            1,
+            1,
+            1,
+            1,
+            1,
+            1,
             0,
             stream.handle,
             params,
-            cooperative=False
+            cooperative=False,
         )
 
     def _ctypes_pointer(self, array):
@@ -158,7 +168,8 @@ class _Runtime:
         self.ensure_allocated()
 
         self._single_thread_launch(
-            self._memsys_module, stream, "NRT_MemSys_init")
+            self._memsys_module, stream, "NRT_MemSys_init"
+        )
         self._initialized = True
 
         if config.CUDA_NRT_STATS:
@@ -170,7 +181,8 @@ class _Runtime:
         Enable memsys statistics
         """
         self._single_thread_launch(
-            self._memsys_module, stream, "NRT_MemSys_enable_stats")
+            self._memsys_module, stream, "NRT_MemSys_enable_stats"
+        )
 
     @_alloc_init_guard
     def memsys_disable_stats(self, stream=None):
@@ -178,7 +190,8 @@ class _Runtime:
         Disable memsys statistics
         """
         self._single_thread_launch(
-            self._memsys_module, stream, "NRT_MemSys_disable_stats")
+            self._memsys_module, stream, "NRT_MemSys_disable_stats"
+        )
 
     @_alloc_init_guard
     def memsys_stats_enabled(self, stream=None):
@@ -193,7 +206,7 @@ class _Runtime:
             self._memsys_module,
             stream,
             "NRT_MemSys_stats_enabled",
-            (enabled_ptr,)
+            (enabled_ptr,),
         )
 
         cuda.synchronize()
@@ -204,21 +217,20 @@ class _Runtime:
         """
         Copy all statistics of memsys to the host
         """
-        dt = np.dtype([
-            ('alloc', np.uint64),
-            ('free', np.uint64),
-            ('mi_alloc', np.uint64),
-            ('mi_free', np.uint64)
-        ])
+        dt = np.dtype(
+            [
+                ("alloc", np.uint64),
+                ("free", np.uint64),
+                ("mi_alloc", np.uint64),
+                ("mi_free", np.uint64),
+            ]
+        )
 
         stats_for_read = cuda.managed_array(1, dt)
         stats_ptr = self._ctypes_pointer(stats_for_read)
 
         self._single_thread_launch(
-            self._memsys_module,
-            stream,
-            "NRT_MemSys_read",
-            [stats_ptr]
+            self._memsys_module, stream, "NRT_MemSys_read", [stats_ptr]
         )
         cuda.synchronize()
 
@@ -237,7 +249,7 @@ class _Runtime:
             alloc=memsys["alloc"],
             free=memsys["free"],
             mi_alloc=memsys["mi_alloc"],
-            mi_free=memsys["mi_free"]
+            mi_free=memsys["mi_free"],
         )
 
     @_alloc_init_guard
@@ -249,10 +261,7 @@ class _Runtime:
         got_ptr = self._ctypes_pointer(got)
 
         self._single_thread_launch(
-            self._memsys_module,
-            stream,
-            f"NRT_MemSys_read_{stat}",
-            [got_ptr]
+            self._memsys_module, stream, f"NRT_MemSys_read_{stat}", [got_ptr]
         )
 
         cuda.synchronize()
@@ -309,15 +318,13 @@ class _Runtime:
         """
         if self._memsys is None:
             raise RuntimeError(
-                "Please allocate NRT Memsys first before setting to module.")
+                "Please allocate NRT Memsys first before setting to module."
+            )
 
         memsys_ptr = self._ctypes_pointer(self._memsys)
 
         self._single_thread_launch(
-            module,
-            stream,
-            "NRT_MemSys_set",
-            [memsys_ptr]
+            module, stream, "NRT_MemSys_set", [memsys_ptr]
         )
 
     @_alloc_init_guard
@@ -327,9 +334,7 @@ class _Runtime:
         """
         cuda.synchronize()
         self._single_thread_launch(
-            self._memsys_module,
-            stream,
-            "NRT_MemSys_print"
+            self._memsys_module, stream, "NRT_MemSys_print"
         )
 
 

numba_cuda/numba/cuda/simulator/__init__.py

@@ -3,14 +3,22 @@ import sys
 from .api import *
 from .vector_types import vector_types
 from .reduction import Reduce
-from .cudadrv.devicearray import (device_array, device_array_like, pinned,
-                                  pinned_array, pinned_array_like,
-                                  mapped_array, to_device, auto_device)
+from .cudadrv.devicearray import (
+    device_array,
+    device_array_like,
+    pinned,
+    pinned_array,
+    pinned_array_like,
+    mapped_array,
+    to_device,
+    auto_device,
+)
 from .cudadrv import devicearray
 from .cudadrv.devices import require_context, gpus
 from .cudadrv.devices import get_context as current_context
 from .cudadrv.runtime import runtime
 from numba.core import config
+
 reduce = Reduce
 
 # Register simulated vector types as module level variables
@@ -25,14 +33,16 @@ del vector_types, name, svty, alias
 if config.ENABLE_CUDASIM:
     import sys
     from numba.cuda.simulator import cudadrv
-    sys.modules['numba.cuda.cudadrv'] = cudadrv
-    sys.modules['numba.cuda.cudadrv.devicearray'] = cudadrv.devicearray
-    sys.modules['numba.cuda.cudadrv.devices'] = cudadrv.devices
-    sys.modules['numba.cuda.cudadrv.driver'] = cudadrv.driver
-    sys.modules['numba.cuda.cudadrv.runtime'] = cudadrv.runtime
-    sys.modules['numba.cuda.cudadrv.drvapi'] = cudadrv.drvapi
-    sys.modules['numba.cuda.cudadrv.error'] = cudadrv.error
-    sys.modules['numba.cuda.cudadrv.nvvm'] = cudadrv.nvvm
+
+    sys.modules["numba.cuda.cudadrv"] = cudadrv
+    sys.modules["numba.cuda.cudadrv.devicearray"] = cudadrv.devicearray
+    sys.modules["numba.cuda.cudadrv.devices"] = cudadrv.devices
+    sys.modules["numba.cuda.cudadrv.driver"] = cudadrv.driver
+    sys.modules["numba.cuda.cudadrv.runtime"] = cudadrv.runtime
+    sys.modules["numba.cuda.cudadrv.drvapi"] = cudadrv.drvapi
+    sys.modules["numba.cuda.cudadrv.error"] = cudadrv.error
+    sys.modules["numba.cuda.cudadrv.nvvm"] = cudadrv.nvvm
 
     from . import compiler
-    sys.modules['numba.cuda.compiler'] = compiler
+
+    sys.modules["numba.cuda.compiler"] = compiler