Skip to content

fix: chunk multinormal and multinomial draws so each oneMKL call fits in MKL_INT - #187

Open
intel-python-devops wants to merge 3 commits into
masterfrom
chore/agentic-sweep-14
Open

intel-python-devops wants to merge 3 commits into
masterfrom
chore/agentic-sweep-14

Conversation

@intel-python-devops

@intel-python-devops intel-python-devops commented Oct 5, 2026 •

Copy link
Copy Markdown

Problem

irk_multinormal_vec_ICDF passed its npy_intp len straight into the 32-bit MKL_INT count of vdRngGaussianMV, with no chunking loop and no len < 1 guard, unlike irk_multinormal_vec_BM1/BM2.

Chunking at MKL_INT_MAX vectors (what BM1/BM2 did) is still not enough for dim >= 2. oneMKL's 32-bit interface mishandles a single vdRngGaussianMV call once n * dim > MKL_INT_MAX: the call returns VSL_STATUS_OK but generates nothing, or only a wrapped-around prefix. Since np.empty typically returns zeroed pages, the user sees every sample silently equal to mean. This affects ICDF, BoxMuller and BoxMuller2, and starts well below MKL_INT_MAX samples (e.g. at n = 2**30 + 1 for dim = 2).

Changes

  • mkl_random/src/mkl_distributions.cpp
    • irk_multinormal_vec_ICDF, irk_multinormal_vec_BM1 and irk_multinormal_vec_BM2 now return early for len < 1 and split the request into chunks of MKL_INT_MAX / dim vectors, advancing the output pointer by max_len * dim doubles per chunk.
    • irk_multinomial_vec chunks at MKL_INT_MAX / k the same way. This is defensive: it was already correct at the sizes tested.
  • mkl_random/tests/test_random.py: two behaviour tests for multinormal_cholesky.
    • Splitting a draw into two sequential calls gives the same values as one combined call, for ICDF, BoxMuller and BoxMuller2.
    • An empty size (0 and (3, 0)) returns an empty array of the right shape and leaves the stream untouched.
  • CHANGELOG.md: added a Fixed entry.

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 n samples 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):

Case master PR
ICDF, dim=2, n=2**30+1 (16 GiB) every row equals mean 0 mismatches
ICDF, dim=1, n=2**31+1 (16 GiB) oneMKL error, output never written 0 mismatches
BoxMuller / BoxMuller2, dim=2, n=2**30+1 (16 GiB) not run 0 mismatches
ICDF, dim=3, n=1.5e9 (33.5 GiB) partly filled, rest equals mean 0 mismatches
ICDF, dim=2, n=231+220 (32 GiB) oneMKL error, output never written 0 mismatches
multinomial, k=2, n=230+220 (8 GiB) 0 mismatches 0 mismatches

A build with MKL_INT_MAX forced 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.

- 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.
@antonwolfy

Copy link
Copy Markdown
Collaborator

Nit: the subject of f924477, which is also the PR title, is cut off. It ends with "...like its BM1/BM".

Suggested: fix: chunk irk_multinormal_vec_ICDF draws at MKL_INT_MAX like its BM1/BM2 siblings

The repo merges with merge commits, so the truncated subject will land in master as-is unless the commit is reworded before merging.

@antonwolfy antonwolfy added this to the 1.6.0 release milestone Oct 6, 2026
@antonwolfy

Copy link
Copy Markdown
Collaborator

Large-size validation results

The 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 n samples in a single call with MKLRandomState(777), then redraw the same stream from a fresh state with the same seed in pieces of 2**24 samples (each far below any limit) and compare row by row.

Three builds were compared:

  • master
  • this PR as currently pushed (4df72af)
  • this PR plus a proposed follow-up that chunks by elements instead of by vectors (diff below)
Case Output size master PR as pushed PR + follow-up
A ICDF, dim=2, n=2**30+1 16 GiB ❌ every row equals mean ❌ same as master ✅ 0 mismatches
B ICDF, dim=1, n=2**31+1 16 GiB ❌ oneMKL "Parameter 3 was incorrect", output never written ✅ 0 mismatches ✅ 0 mismatches
C BoxMuller, dim=2, n=2**30+1 16 GiB n/a n/a ✅ 0 mismatches
D BoxMuller2, dim=2, n=2**30+1 16 GiB n/a n/a ✅ 0 mismatches
E ICDF, dim=3, n=1.5e9 33.5 GiB ❌ first 68,344,234 rows correct, the rest wrong (all but one row equal mean) ❌ same as master ✅ 0 mismatches
G ICDF, dim=2, n=231+220 32 GiB ❌ oneMKL "Parameter 3" error, output never written ❌ every row wrong (see below) ✅ 0 mismatches
F multinomial, k=2, ntrial=2, n=230+220 8 GiB ✅ 0 mismatches n/a ✅ 0 mismatches

"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

  1. The PR as pushed fixes only dim == 1 (case B). For dim >= 2 it behaves like master, because the bad range starts well below MKL_INT_MAX samples.
  2. oneMKL's 32-bit interface mishandles a single vdRngGaussianMV call once n * dim > 2**31 - 1. The call returns VSL_STATUS_OK but generates nothing, so the output keeps whatever the buffer held. A fresh np.empty is typically zero pages, so every sample silently equals mean. This reproduces with ICDF, BoxMuller and BoxMuller2, and is independent of where the MKL_INT_MAX chunk boundary falls. For example, with dim=2 it already breaks at n = 2**30 + 1.
  3. Case G shows the PR's own loop is not enough for dim >= 2. Its first chunk has MKL_INT_MAX vectors, so chunk * dim overflows: all 2,147,483,647 rows of that chunk equal mean, and the remaining rows come from the wrong stream position.
  4. Case E (partial fill): 4.5e9 elements wrap to 205,032,704 in 32 bits, which is 68,344,234.67 vectors. That matches the observed boundary exactly.
  5. Multinomial (case F) was correct at this size on this oneMKL build, with and without the follow-up. Chunking at MKL_INT_MAX / k there is a defensive change only.

Proposed follow-up

In irk_multinormal_vec_ICDF, irk_multinormal_vec_BM1 and irk_multinormal_vec_BM2, chunk at MKL_INT_MAX / dim vectors instead of MKL_INT_MAX:

    /* 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 irk_multinomial_vec uses MKL_INT_MAX / k, and the CHANGELOG entry is reworded to describe the general condition (n * dim > MKL_INT_MAX). With this applied, all cases above pass. I also ran the normal and a forced-small-chunk build (MKL_INT_MAX set to 7) against each other over 1,480 combinations of dimension, storage layout, method and sample count, and they agree to within last-bit rounding (max 8.9e-16).

Note on the unit tests

The 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.

Reproducer
import 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 0

Needs 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.
@antonwolfy

Copy link
Copy Markdown
Collaborator

The proposed follow-up from the previous comment is now on this branch as 42a8f94:

  • irk_multinormal_vec_ICDF, irk_multinormal_vec_BM1 and irk_multinormal_vec_BM2 chunk at MKL_INT_MAX / dim vectors, and advance the output pointer by max_len * dim doubles.
  • irk_multinomial_vec chunks at MKL_INT_MAX / k (defensive; case F was already correct).
  • The split-call test now runs for ICDF, BoxMuller and BoxMuller2.
  • The CHANGELOG entry now describes the general n * dim > MKL_INT_MAX condition.

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.

@antonwolfy antonwolfy changed the title fix: chunk irk_multinormal_vec_ICDF draws at MKL_INT_MAX like its BM1/BM fix: chunk multinormal and multinomial draws so each oneMKL call fits in MKL_INT Oct 6, 2026

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants