From e76b6b6f672558d9f56f2374821938d481c06c98 Mon Sep 17 00:00:00 2001 From: Vladislav Perevezentsev Date: Fri, 4 Sep 2026 06:22:48 -0700 Subject: [PATCH 1/4] Speed up randint array-bounds via chunking Lemire --- mkl_random/src/mkl_distributions.cpp | 113 +++++++++++++++------------ 1 file changed, 63 insertions(+), 50 deletions(-) diff --git a/mkl_random/src/mkl_distributions.cpp b/mkl_random/src/mkl_distributions.cpp index 01cd256..b9cb7c3 100644 --- a/mkl_random/src/mkl_distributions.cpp +++ b/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) From 01547a2976166d39ea98bb411836410579f530ba Mon Sep 17 00:00:00 2001 From: Vladislav Perevezentsev Date: Fri, 4 Sep 2026 06:35:48 -0700 Subject: [PATCH 2/4] Add gh-173 to changelog --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 61e54d3..79be316 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Changed * Pinned Cython in the Coverity Scan workflow so generated code stays stable between scans, and added `coverity/README.md` documenting the known Cython-boilerplate false positives and the scan review checklist [gh-164](https://github.com/IntelPython/mkl_random/pull/164) +* Speed up `randint` with `array_like` bounds via word chunking and a branchless Lemire loop [gh-173](https://github.com/IntelPython/mkl_random/pull/173) ### Fixed * Fixed `uniform` to return a Python `float` for scalar bounds with `size=None` instead of a 0-d array [gh-167](https://github.com/IntelPython/mkl_random/pull/167) From b733136800567b1a7c0f0884685434f4f19b9a2c Mon Sep 17 00:00:00 2001 From: Vladislav Perevezentsev Date: Wed, 16 Sep 2026 05:24:36 -0700 Subject: [PATCH 3/4] Drop the branchless Lemire tweak from randint broadcast --- mkl_random/src/mkl_distributions.cpp | 40 ++++++++++++---------------- 1 file changed, 17 insertions(+), 23 deletions(-) diff --git a/mkl_random/src/mkl_distributions.cpp b/mkl_random/src/mkl_distributions.cpp index 8e7128c..05e9083 100644 --- a/mkl_random/src/mkl_distributions.cpp +++ b/mkl_random/src/mkl_distributions.cpp @@ -2289,9 +2289,6 @@ static void irk_rand_bounded_broadcast(irk_state *state, WT *words = (WT *)mkl_malloc(chunk_cap * sizeof(WT), 64); assert(words != nullptr); - /* 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; @@ -2310,19 +2307,17 @@ static void irk_rand_bounded_broadcast(irk_state *state, if (s != 0) { WT lo = 0; result = irk_mulhi(w, s, &lo); - /* 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); + if (lo < s) { /* rare */ + WT t = (WT)(0 - s) % s; + 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; } - idx[n_pending++] = j; - continue; } } res[j] = (T)(((UT)low[j]) + (UT)result); @@ -2344,14 +2339,13 @@ static void irk_rand_bounded_broadcast(irk_state *state, 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; + 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); From 5c2319a480ee88d816ddf58ffc2e2a2ffd1e5fb8 Mon Sep 17 00:00:00 2001 From: Vladislav Perevezentsev Date: Wed, 16 Sep 2026 06:03:05 -0700 Subject: [PATCH 4/4] Pick the Lemire rejection test by range width --- mkl_random/src/mkl_distributions.cpp | 75 +++++++++++++++++++++------- 1 file changed, 58 insertions(+), 17 deletions(-) diff --git a/mkl_random/src/mkl_distributions.cpp b/mkl_random/src/mkl_distributions.cpp index 05e9083..d7b2211 100644 --- a/mkl_random/src/mkl_distributions.cpp +++ b/mkl_random/src/mkl_distributions.cpp @@ -2263,7 +2263,7 @@ 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 by MKL in cache-sized chunks; the rare rejected + * 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, @@ -2289,27 +2289,35 @@ static void irk_rand_bounded_broadcast(irk_state *state, WT *words = (WT *)mkl_malloc(chunk_cap * sizeof(WT), 64); assert(words != 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; + 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); - 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; - if (lo < t) { + if (wide) { + WT last_s = 0, last_t = 0; /* memoized reject threshold */ + + 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); @@ -2319,8 +2327,41 @@ static void irk_rand_bounded_broadcast(irk_state *state, continue; } } + res[j] = (T)(((UT)low[j]) + (UT)result); } - res[j] = (T)(((UT)low[j]) + (UT)result); + } + 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); + } + + if (base == 0) + wide = n_hits > chunk / 16; } /* retry the chunk's rejects locally with fresh words */