Speed up randint array-bounds via chunking Lemire

Author: vlad-perevezentsev

mkl_random/src/mkl_distributions.cpp

@@ -2262,8 +2262,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 rare 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).
  */
@@ -2274,76 +2275,88 @@ 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);
+    /* 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;
+
+    WT *words = (WT *)mkl_malloc(chunk_cap * sizeof(WT), 64);
     assert(words != nullptr);
 
-    irk_uniform_bits_vec(state, len, words);
+    /* memoized reject threshold */
+    WT last_s = 0, last_t = 0;
 
-    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);
-                    }
-                    idx[n_pending++] = i;
-                    continue;
-                }
-            }
-        }
-        res[i] = (T)(((UT)low[i]) + (UT)result);
-    }
-
-    while (n_pending > 0) {
-        npy_intp wpos = 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, n_pending, words);
+        irk_uniform_bits_vec(state, chunk, words);
 
-        for (k = 0; k < n_pending; ++k) {
-            npy_intp j = idx[k];
-            WT w = (WT)words[k];
+        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;
+            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) {
-                    WT t = (WT)(0 - s) % s;
-                    if (lo < t) {
+                /* t < s always (no lo < s branch) */
+                if (s != last_s) { /* recompute threshold */
+                    last_t = (WT)(0 - s) % s;
+                    last_s = s;
+                }
+                if (lo < last_t) { /* rare reject */
+                    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);
+        }
+
+        /* 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 (s != last_s) {
+                        last_t = (WT)(0 - s) % s;
+                        last_s = s;
+                    }
+                    if (lo < last_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)