Repository navigation
fix: chunk multinormal and multinomial draws so each oneMKL call fits in MKL_INT - #187
intel-python-devops wants to merge 3 commits into
Conversation
- Shorten the split-call test comment and compare with assert_allclose, since bit-exact continuation across calls is an MKL implementation detail that does not hold for every BRNG. - Reword the empty-size test: a 0-count MKL call was already a no-op, so check instead that an empty draw returns an empty array and leaves the stream untouched, for all three methods and for size=0 and (3, 0). - Add a Fixed entry to CHANGELOG.md.
|
Nit: the subject of f924477, which is also the PR title, is cut off. It ends with "...like its BM1/BM". Suggested: The repo merges with merge commits, so the truncated subject will land in |
Large-size validation resultsThe new unit tests only exercise small sizes, so I ran the multivariate-normal and multinomial paths at sizes past the 32-bit limit (up to ~33 GiB outputs). Setup: Linux x86_64, Intel Xeon (Sapphire Rapids, AVX-512), oneMKL 2026.2.0 (LP64), Python 3.13, NumPy 2.5.3, GCC 15.3. Method: draw Three builds were compared:
"0 mismatches" means bit-for-bit equal to the piecewise reference. The unit-test suite also passes on all three builds (338 / 345 / 347 passed, 24 skipped, 0 failed). Findings
Proposed follow-upIn /* oneMKL fills len * dim doubles per call, and that count must fit in
MKL_INT. */
const npy_intp max_len = MKL_INT_MAX / dim;
while (len > max_len) {
err = vdRngGaussianMV(VSL_RNG_METHOD_GAUSSIANMV_ICDF, state->stream,
max_len, res, dim, storage_mode, mean_vec, ch);
assert(err == VSL_STATUS_OK);
res += max_len * dim;
len -= max_len;
}The same change in Note on the unit testsThe two small-size tests in this PR also pass on master, so they check behaviour rather than this fix. A true regression test needs at least 16 GiB of output, which is why the large-size checks above were run by hand. Reproducerimport numpy as np
import mkl_random as rnd
n, dim, step = 2**30 + 1, 2, 2**24
mean = np.array([0.5, -0.25])
ch = np.array([[1.0, 0.0], [0.3, 1.0]])
r = rnd.MKLRandomState(777).multinormal_cholesky(mean, ch, size=n, method="ICDF")
s = rnd.MKLRandomState(777)
bad = 0
for i in range(0, n, step):
m = min(step, n - i)
ref = s.multinormal_cholesky(mean, ch, size=m, method="ICDF")
bad += int(np.count_nonzero(np.any(r[i:i + m] != ref, axis=1)))
print("rows differing from piecewise reference:", bad) # expected 0Needs about 17 GiB of free memory. |
… in MKL_INT oneMKL's 32-bit interface mishandles a single vdRngGaussianMV call once n * dim exceeds MKL_INT_MAX: it returns VSL_STATUS_OK but generates nothing (or only a wrapped-around prefix). Chunking at MKL_INT_MAX vectors, as irk_multinormal_vec_BM1/BM2 did, is therefore not enough for dim >= 2. - Chunk irk_multinormal_vec_ICDF/BM1/BM2 at MKL_INT_MAX / dim vectors and advance the output pointer by max_len * dim doubles. - Chunk irk_multinomial_vec at MKL_INT_MAX / k as a defensive measure. - Run the split-call test for all three methods. - Reword the CHANGELOG entry to cover the general n * dim condition.
|
The proposed follow-up from the previous comment is now on this branch as 42a8f94:
The large-size results above (cases A–G) were measured with this change applied. The PR title and description still describe only the original ICDF fix; I'll update them separately. |
Problem
irk_multinormal_vec_ICDFpassed itsnpy_intp lenstraight into the 32-bitMKL_INTcount ofvdRngGaussianMV, with no chunking loop and nolen < 1guard, unlikeirk_multinormal_vec_BM1/BM2.Chunking at
MKL_INT_MAXvectors (what BM1/BM2 did) is still not enough fordim >= 2. oneMKL's 32-bit interface mishandles a singlevdRngGaussianMVcall oncen * dim > MKL_INT_MAX: the call returnsVSL_STATUS_OKbut generates nothing, or only a wrapped-around prefix. Sincenp.emptytypically returns zeroed pages, the user sees every sample silently equal tomean. This affectsICDF,BoxMullerandBoxMuller2, and starts well belowMKL_INT_MAXsamples (e.g. atn = 2**30 + 1fordim = 2).Changes
mkl_random/src/mkl_distributions.cppirk_multinormal_vec_ICDF,irk_multinormal_vec_BM1andirk_multinormal_vec_BM2now return early forlen < 1and split the request into chunks ofMKL_INT_MAX / dimvectors, advancing the output pointer bymax_len * dimdoubles per chunk.irk_multinomial_vecchunks atMKL_INT_MAX / kthe same way. This is defensive: it was already correct at the sizes tested.mkl_random/tests/test_random.py: two behaviour tests formultinormal_cholesky.ICDF,BoxMullerandBoxMuller2.size(0and(3, 0)) returns an empty array of the right shape and leaves the stream untouched.CHANGELOG.md: added aFixedentry.Testing
The two unit tests above also pass on master, so they are behaviour tests, not regression tests: reproducing the bug needs an output of at least 16 GiB, which CI cannot allocate.
The large-size behaviour was checked by hand instead. Each case draws
nsamples in one call, redraws them from a fresh state with the same seed in pieces of 2**24 samples, and compares row by row. Full results and a reproducer are in the comments below. Summary (oneMKL 2026.2.0, LP64):ICDF, dim=2, n=2**30+1 (16 GiB)meanICDF, dim=1, n=2**31+1 (16 GiB)BoxMuller/BoxMuller2, dim=2, n=2**30+1 (16 GiB)ICDF, dim=3, n=1.5e9 (33.5 GiB)meanICDF, dim=2, n=231+220 (32 GiB)A build with
MKL_INT_MAXforced to 7 also agrees with the normal build over 1,480 combinations of dimension, storage layout, method and sample count, to within last-bit rounding.