Speed up randint array-bounds via chunking Lemire

CHANGELOG.md

@@ -22,6 +22,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 * Raised the minimum build-time `Cython` requirement to `3.1.0`, the first release providing the `freethreading_compatible` directive [gh-159](https://github.com/IntelPython/mkl_random/pull/159)
 * Extended the `memcpy`-based fast path of `shuffle` to multi-dimensional `ndarray` inputs whose first-axis items are contiguous, which is also much faster than the previous buffered path [gh-159](https://github.com/IntelPython/mkl_random/pull/159)
 * Updated the `randint_untyped` implementation, which now supports `array_like` (broadcastable) `low`/`high` bounds [gh-176](https://github.com/IntelPython/mkl_random/pull/176)
+* Sped up `randint` with `array_like` bounds by generating words in cache-sized chunks and by picking the Lemire rejection test by range width [gh-173](https://github.com/IntelPython/mkl_random/pull/173)
 
 ### Fixed
 * Fixed an out-of-range integer `brng` indexing `brng_list` past its end, which returned uninitialized memory as random values; it now warns and falls back to `MT19937` [gh-177](https://github.com/IntelPython/mkl_random/pull/177)

mkl_random/src/mkl_distributions.cpp

@@ -2054,8 +2054,9 @@ static inline npy_uint64 irk_mulhi(npy_uint64 a, npy_uint64 b, npy_uint64 *lo)
 /*
  * Draw res[i] uniformly from [low[i], hi[i]] (inclusive) using Lemire's
  * multiply-shift method (per-element bounds, same as NumPy).
- * Words are generated in bulk by MKL; the rare rejected elements are
- * gathered into `idx` (allocated lazily) and retried on the next round.
+ * Words are generated by MKL in cache-sized chunks; the rejected
+ * elements are gathered into `idx` (allocated lazily) and retried
+ * locally within each chunk.
  * T is the result type, UT its unsigned counterpart,
  * WT the raw-word type (s wraps to 0 for a full-range draw).
  */
@@ -2066,76 +2067,123 @@ static void irk_rand_bounded_broadcast(irk_state *state,
                                        const T *low,
                                        const T *hi)
 {
-    npy_intp i = 0;
-    npy_intp k = 0;
-    npy_intp n_pending = 0;
-    npy_intp *idx = nullptr;
-    WT *words = nullptr;
+    npy_intp *idx = nullptr; /* reject indices */
 
     if (len < 1)
         return;
 
-    /* TODO: possible speedup :
-     * generate and consume words in cache-sized chunks
-     * instead of one full-length pass */
-    words = (WT *)mkl_malloc(len * sizeof(WT), 64);
-    assert(words != nullptr);
+    /* Optimized path:
+     * cache-sized chunks instead of one full-length pass */
+    const npy_intp CHUNK_SIZE = 1 << 15; /* ~32K elements per chunk */
+    npy_intp chunk_cap = (len < CHUNK_SIZE) ? len : CHUNK_SIZE;
 
-    irk_uniform_bits_vec(state, len, words);
+    WT *words = (WT *)mkl_malloc(chunk_cap * sizeof(WT), 64);
+    assert(words != nullptr);
 
-    for (i = 0; i < len; ++i) {
-        WT w = (WT)words[i];
-        /* diff cast back to UT so narrow types wrap (no signed promotion) */
-        UT d = (UT)(((UT)hi[i]) - ((UT)low[i]));
-        WT s = (WT)d + 1; /* 0 iff full range (32/64-bit only) */
-        WT result = w;
-
-        if (s != 0) {
-            WT lo = 0;
-            result = irk_mulhi(w, s, &lo);
-            if (lo < s) { /* rare */
-                WT t = (WT)(0 - s) % s;
-                if (lo < t) {
-                    if (idx == nullptr) {
-                        idx =
-                            (npy_intp *)mkl_malloc(len * sizeof(npy_intp), 64);
-                        assert(idx != nullptr);
+    /* `lo < s` is free for narrow ranges, but mispredicts for wide ones:
+     * count it on the first chunk, then pick the cheaper test */
+    npy_intp n_hits = 0;
+    bool wide = false;
+    /* memoized reject threshold */
+    WT last_s = 0, last_t = 0;
+
+    for (npy_intp base = 0; base < len; base += chunk_cap) {
+        npy_intp chunk = (len - base < chunk_cap) ? (len - base) : chunk_cap;
+        npy_intp n_pending = 0;
+
+        irk_uniform_bits_vec(state, chunk, words);
+
+        if (wide) {
+            for (npy_intp i = 0; i < chunk; ++i) {
+                npy_intp j = base + i;
+                WT w = (WT)words[i];
+                UT d = (UT)(((UT)hi[j]) - ((UT)low[j]));
+                WT s = (WT)d + 1;
+                WT result = w;
+
+                if (s != 0) {
+                    WT lo = 0;
+                    result = irk_mulhi(w, s, &lo);
+                    if (s != last_s) {
+                        last_t = (WT)(0 - s) % s;
+                        last_s = s;
+                    }
+                    if (lo < last_t) {
+                        if (idx == nullptr) {
+                            idx = (npy_intp *)mkl_malloc(
+                                chunk_cap * sizeof(npy_intp), 64);
+                            assert(idx != nullptr);
+                        }
+                        idx[n_pending++] = j;
+                        continue;
                     }
-                    idx[n_pending++] = i;
-                    continue;
                 }
+                res[j] = (T)(((UT)low[j]) + (UT)result);
             }
         }
-        res[i] = (T)(((UT)low[i]) + (UT)result);
-    }
-
-    while (n_pending > 0) {
-        npy_intp wpos = 0;
-
-        irk_uniform_bits_vec(state, n_pending, words);
+        else {
+            for (npy_intp i = 0; i < chunk; ++i) {
+                npy_intp j = base + i;
+                WT w = (WT)words[i];
+                /* diff cast back to UT so narrow types wrap (no signed
+                 * promotion) */
+                UT d = (UT)(((UT)hi[j]) - ((UT)low[j]));
+                WT s = (WT)d + 1; /* 0 iff full range (32/64-bit only) */
+                WT result = w;
+
+                if (s != 0) {
+                    WT lo = 0;
+                    result = irk_mulhi(w, s, &lo);
+                    if (lo < s) { /* rare */
+                        WT t = (WT)(0 - s) % s;
+                        ++n_hits;
+                        if (lo < t) {
+                            if (idx == nullptr) {
+                                idx = (npy_intp *)mkl_malloc(
+                                    chunk_cap * sizeof(npy_intp), 64);
+                                assert(idx != nullptr);
+                            }
+                            idx[n_pending++] = j;
+                            continue;
+                        }
+                    }
+                }
+                res[j] = (T)(((UT)low[j]) + (UT)result);
+            }
 
-        for (k = 0; k < n_pending; ++k) {
-            npy_intp j = idx[k];
-            WT w = (WT)words[k];
-            UT d = (UT)(((UT)hi[j]) - ((UT)low[j]));
-            WT s = (WT)d + 1;
-            WT result = w;
+            if (base == 0)
+                wide = n_hits > chunk / 16;
+        }
 
-            if (s != 0) {
-                WT lo = 0;
-                result = irk_mulhi(w, s, &lo);
-                if (lo < s) {
-                    WT t = (WT)(0 - s) % s;
-                    if (lo < t) {
-                        /* keep pending; wpos <= k so idx[k] read first */
-                        idx[wpos++] = j;
-                        continue;
+        /* retry the chunk's rejects locally with fresh words */
+        while (n_pending > 0) {
+            npy_intp wpos = 0;
+
+            irk_uniform_bits_vec(state, n_pending, words);
+
+            for (npy_intp k = 0; k < n_pending; ++k) {
+                npy_intp j = idx[k];
+                WT w = (WT)words[k];
+                UT d = (UT)(((UT)hi[j]) - ((UT)low[j]));
+                WT s = (WT)d + 1;
+                WT result = w;
+
+                if (s != 0) {
+                    WT lo = 0;
+                    result = irk_mulhi(w, s, &lo);
+                    if (lo < s) {
+                        WT t = (WT)(0 - s) % s;
+                        if (lo < t) {
+                            /* keep pending; wpos <= k so idx[k] read first */
+                            idx[wpos++] = j;
+                            continue;
+                        }
                     }
                 }
+                res[j] = (T)(((UT)low[j]) + (UT)result);
             }
-            res[j] = (T)(((UT)low[j]) + (UT)result);
+            n_pending = wpos;
         }
-        n_pending = wpos;
     }
 
     if (idx != nullptr)