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

numba_cuda/numba/cuda/random.py

@@ -0,0 +1,308 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import math
+
+from numba import cuda
+from numba.cuda import (
+    float32,
+    float64,
+    uint32,
+    int64,
+    uint64,
+    HAS_NUMBA,
+)
+from numba.cuda.np.numpy_support import from_dtype
+from numba.cuda import config
+
+if HAS_NUMBA:
+    from numba import jit
+
+import numpy as np
+
+# This implementation is based upon the xoroshiro128+ and splitmix64 algorithms
+# described at:
+#
+#     http://xoroshiro.di.unimi.it/
+#
+# and originally implemented by David Blackman and Sebastiano Vigna.
+#
+# The implementations below are based on the C source code:
+#
+#  * http://xoroshiro.di.unimi.it/xoroshiro128plus.c
+#  * http://xoroshiro.di.unimi.it/splitmix64.c
+#
+# Splitmix64 is used to generate the initial state of the xoroshiro128+
+# generator to ensure that small seeds don't result in predictable output.
+
+# **WARNING**: There is a lot of verbose casting in this file to ensure that
+# NumPy casting conventions (which cast uint64 [op] int32 to float64) don't
+# turn integers into floats when using these functions in the CUDA simulator.
+#
+# There are also no function type signatures to ensure that compilation is
+# deferred so that import is quick, and Sphinx autodoc works.  We are also
+# 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_type = from_dtype(xoroshiro128p_dtype)
+
+# When cudasim is enabled, Fake CUDA arrays are passed to some of the
+# @jit-decorated functions. This required fallback to object mode. With
+# Numba 0.59.0 object mode must be explicitly enabled.
+# https://numba.readthedocs.io/en/stable/reference/deprecation.html#deprecation-of-object-mode-fall-back-behaviour-when-using-jit
+# In order to avoid the warning / future error, we explicitly specify that
+# object mode with loop lifting is acceptable when using the simulator.
+_forceobj = _looplift = config.ENABLE_CUDASIM
+_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.
+
+    This ensures that manually set small seeds don't result in a predictable
+    initial sequence from the random number generator.
+
+    :type states: 1D array, dtype=xoroshiro128p_dtype
+    :param states: array of RNG states
+    :type index: uint64
+    :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)
+
+    z = seed + uint64(0x9E3779B97F4A7C15)
+    z = (z ^ (z >> uint32(30))) * uint64(0xBF58476D1CE4E5B9)
+    z = (z ^ (z >> uint32(27))) * uint64(0x94D049BB133111EB)
+    z = z ^ (z >> uint32(31))
+
+    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."""
+    x = uint64(x)
+    k = uint32(k)
+    return (x << k) | (x >> uint32(64 - 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].
+
+    :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"]
+    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)))
+
+    return result
+
+
+@jit(forceobj=_forceobj, looplift=_looplift, nopython=_nopython)
+def xoroshiro128p_jump(states, index):
+    """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))
+
+    s0 = uint64(0)
+    s1 = uint64(0)
+
+    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"]
+            xoroshiro128p_next(states, index)
+
+    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)"""
+    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)"""
+    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]``.
+
+    :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]``.
+
+    :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))
+
+
+TWO_PI_FLOAT32 = np.float32(2 * math.pi)
+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]``.
+
+    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.
+
+    :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)
+
+    u1 = xoroshiro128p_uniform_float32(states, index)
+    u2 = xoroshiro128p_uniform_float32(states, index)
+
+    z0 = math.sqrt(-float32(2.0) * math.log(u1)) * math.cos(TWO_PI_FLOAT32 * u2)
+    # discarding second normal value
+    # z1 = math.sqrt(-float32(2.0) * math.log(u1))
+    #                * math.sin(TWO_PI_FLOAT32 * u2)
+    return z0
+
+
+@jit(forceobj=_forceobj, looplift=_looplift, nopython=_nopython)
+def xoroshiro128p_normal_float64(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.
+
+    :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)
+
+    u1 = xoroshiro128p_uniform_float32(states, index)
+    u2 = xoroshiro128p_uniform_float32(states, index)
+
+    z0 = math.sqrt(-float64(2.0) * math.log(u1)) * math.cos(TWO_PI_FLOAT64 * u2)
+    # discarding second normal value
+    # z1 = math.sqrt(-float64(2.0) * math.log(u1))
+    #                * math.sin(TWO_PI_FLOAT64 * u2)
+    return z0
+
+
+@jit(forceobj=_forceobj, looplift=_looplift, nopython=_nopython)
+def init_xoroshiro128p_states_cpu(states, seed, subsequence_start):
+    n = states.shape[0]
+    seed = uint64(seed)
+    subsequence_start = uint64(subsequence_start)
+
+    if n >= 1:
+        init_xoroshiro128p_state(states, 0, seed)
+
+        # advance to starting subsequence number
+        for _ in range(subsequence_start):
+            xoroshiro128p_jump(states, 0)
+
+        # populate the rest of the array
+        for i in range(1, n):
+            states[i] = states[i - 1]  # take state of previous generator
+            xoroshiro128p_jump(states, i)  # and jump forward 2**64 steps
+
+
+def init_xoroshiro128p_states(states, seed, subsequence_start=0, stream=0):
+    """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
+    sequence.  Therefore, as long no CUDA thread requests more than 2**64
+    random numbers, all of the RNG states produced by this function are
+    guaranteed to be independent.
+
+    The subsequence_start parameter can be used to advance the first RNG state
+    by a multiple of 2**64 steps.
+
+    :type states: 1D DeviceNDArray, dtype=xoroshiro128p_dtype
+    :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)
+    init_xoroshiro128p_states_cpu(states_cpu, seed, subsequence_start)
+
+    states.copy_to_device(states_cpu, stream=stream)
+
+
+def create_xoroshiro128p_states(n, seed, subsequence_start=0, stream=0):
+    """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
+    sequence.  Therefore, as long no CUDA thread requests more than 2**64
+    random numbers, all of the RNG states produced by this function are
+    guaranteed to be independent.
+
+    The subsequence_start parameter can be used to advance the first RNG state
+    by a multiple of 2**64 steps.
+
+    :type n: int
+    :param n: number of RNG states to create
+    :type seed: uint64
+    :param seed: starting seed for list of generators
+    :type subsequence_start: uint64
+    :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

@@ -0,0 +1,156 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-2-Clause
+ */
+
+/*
+ * Handle reshaping of zero-sized array.
+ * See numba_attempt_nocopy_reshape() below.
+ */
+#define NPY_MAXDIMS 32
+
+typedef long long int npy_intp;
+
+extern "C" __device__ int
+nocopy_empty_reshape(npy_intp nd, const npy_intp *dims, const npy_intp *strides,
+                     npy_intp newnd, const npy_intp *newdims,
+                     npy_intp *newstrides, npy_intp itemsize,
+                     int is_f_order)
+{
+    int i;
+    /* Just make the strides vaguely reasonable
+     * (they can have any value in theory).
+     */
+    for (i = 0; i < newnd; i++)
+        newstrides[i] = itemsize;
+    return 1;  /* reshape successful */
+}
+
+/*
+ * Straight from Numpy's _attempt_nocopy_reshape()
+ * (np/core/src/multiarray/shape.c).
+ * Attempt to reshape an array without copying data
+ *
+ * This function should correctly handle all reshapes, including
+ * axes of length 1. Zero strides should work but are untested.
+ *
+ * If a copy is needed, returns 0
+ * If no copy is needed, returns 1 and fills `npy_intp *newstrides`
+ *     with appropriate strides
+ */
+extern "C" __device__ int
+numba_attempt_nocopy_reshape(npy_intp nd, const npy_intp *dims, const npy_intp *strides,
+                             npy_intp newnd, const npy_intp *newdims,
+                             npy_intp *newstrides, npy_intp itemsize,
+                             int is_f_order)
+{
+    int oldnd;
+    npy_intp olddims[NPY_MAXDIMS];
+    npy_intp oldstrides[NPY_MAXDIMS];
+    npy_intp np, op, last_stride;
+    int oi, oj, ok, ni, nj, nk;
+
+    oldnd = 0;
+    /*
+     * Remove axes with dimension 1 from the old array. They have no effect
+     * but would need special cases since their strides do not matter.
+     */
+    for (oi = 0; oi < nd; oi++) {
+        if (dims[oi]!= 1) {
+            olddims[oldnd] = dims[oi];
+            oldstrides[oldnd] = strides[oi];
+            oldnd++;
+        }
+    }
+
+    np = 1;
+    for (ni = 0; ni < newnd; ni++) {
+        np *= newdims[ni];
+    }
+    op = 1;
+    for (oi = 0; oi < oldnd; oi++) {
+        op *= olddims[oi];
+    }
+    if (np != op) {
+        /* different total sizes; no hope */
+        return 0;
+    }
+
+    if (np == 0) {
+        /* the Numpy code does not handle 0-sized arrays */
+        return nocopy_empty_reshape(nd, dims, strides,
+                                    newnd, newdims, newstrides,
+                                    itemsize, is_f_order);
+    }
+
+    /* oi to oj and ni to nj give the axis ranges currently worked with */
+    oi = 0;
+    oj = 1;
+    ni = 0;
+    nj = 1;
+    while (ni < newnd && oi < oldnd) {
+        np = newdims[ni];
+        op = olddims[oi];
+
+        while (np != op) {
+            if (np < op) {
+                /* Misses trailing 1s, these are handled later */
+                np *= newdims[nj++];
+            } else {
+                op *= olddims[oj++];
+            }
+        }
+
+        /* Check whether the original axes can be combined */
+        for (ok = oi; ok < oj - 1; ok++) {
+            if (is_f_order) {
+                if (oldstrides[ok+1] != olddims[ok]*oldstrides[ok]) {
+                     /* not contiguous enough */
+                    return 0;
+                }
+            }
+            else {
+                /* C order */
+                if (oldstrides[ok] != olddims[ok+1]*oldstrides[ok+1]) {
+                    /* not contiguous enough */
+                    return 0;
+                }
+            }
+        }
+
+        /* Calculate new strides for all axes currently worked with */
+        if (is_f_order) {
+            newstrides[ni] = oldstrides[oi];
+            for (nk = ni + 1; nk < nj; nk++) {
+                newstrides[nk] = newstrides[nk - 1]*newdims[nk - 1];
+            }
+        }
+        else {
+            /* C order */
+            newstrides[nj - 1] = oldstrides[oj - 1];
+            for (nk = nj - 1; nk > ni; nk--) {
+                newstrides[nk - 1] = newstrides[nk]*newdims[nk];
+            }
+        }
+        ni = nj++;
+        oi = oj++;
+    }
+
+    /*
+     * Set strides corresponding to trailing 1s of the new shape.
+     */
+    if (ni >= 1) {
+        last_stride = newstrides[ni - 1];
+    }
+    else {
+        last_stride = itemsize;
+    }
+    if (is_f_order) {
+        last_stride *= newdims[ni - 1];
+    }
+    for (nk = ni; nk < newnd; nk++) {
+        newstrides[nk] = last_stride;
+    }
+
+    return 1;
+}