diff --git a/CHANGELOG.md b/CHANGELOG.md index 265c3b2..d441b70 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -15,6 +15,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 * 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) * 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) +* 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) diff --git a/mkl_random/src/mkl_distributions.cpp b/mkl_random/src/mkl_distributions.cpp index 5eb6071..d7b2211 100644 --- a/mkl_random/src/mkl_distributions.cpp +++ b/mkl_random/src/mkl_distributions.cpp @@ -2263,8 +2263,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). */ @@ -2275,76 +2276,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; + + 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) { + 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); + 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)