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

CHANGELOG.md

@@ -39,6 +39,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 * Fixed `randint` with scalar bounds above `INT_MAX`, `tomaxint` and `bytes` returning uninitialized data for `WH`, `MCG31`, `R250` and `MRG32K3A`, which lack `viRngUniformBits` support [gh-175](https://github.com/IntelPython/mkl_random/pull/175)
 * Fixed 64-bit integer generation returning zeros or crashing with `PHILOX4X32X10` and `ARS5` for requests of `2**30` elements or more [gh-175](https://github.com/IntelPython/mkl_random/pull/175)
 * Fixed `randint_untyped` raising `OverflowError` for bounds outside the C `long` range, e.g. `2**40` on Windows [gh-176](https://github.com/IntelPython/mkl_random/pull/176)
+* Fixed `multinormal_cholesky` silently returning wrong output, e.g. every sample equal to `mean`, when the number of samples times the dimension exceeds `MKL_INT_MAX` [gh-187](https://github.com/IntelPython/mkl_random/pull/187)
 
 ### Removed
 * Removed the `python-gil` constraint from the conda recipes, which pinned `mkl_random` to GIL-enabled Python 3.14 builds [gh-159](https://github.com/IntelPython/mkl_random/pull/159)

mkl_random/src/mkl_distributions.cpp

@@ -1284,13 +1284,17 @@ void irk_multinomial_vec(irk_state *state,
         memset(res, 0, len * k * sizeof(int));
     }
     else {
-        while (len > MKL_INT_MAX) {
+        /* Keep len * k within MKL_INT, since oneMKL indexes the output
+           in MKL_INT. */
+        const npy_intp max_len = MKL_INT_MAX / k;
+
+        while (len > max_len) {
             err = viRngMultinomial(VSL_RNG_METHOD_MULTINOMIAL_MULTPOISSON,
-                                   state->stream, MKL_INT_MAX, res, n, k, pvec);
+                                   state->stream, max_len, res, n, k, pvec);
             assert(err == VSL_STATUS_OK);
             /* len counts draws, res counts ints. */
-            res += k * MKL_INT_MAX;
-            len -= MKL_INT_MAX;
+            res += k * max_len;
+            len -= max_len;
         }
 
         err = viRngMultinomial(VSL_RNG_METHOD_MULTINOMIAL_MULTPOISSON,
@@ -2287,6 +2291,22 @@ void irk_multinormal_vec_ICDF(irk_state *state,
     int err = 0;
     const MKL_INT storage_mode = cholesky_storage_flags[storage_flag];
 
+    if (len < 1)
+        return;
+
+    /* 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;
+    }
+
     err = vdRngGaussianMV(VSL_RNG_METHOD_GAUSSIANMV_ICDF, state->stream, len,
                           res, dim, storage_mode, mean_vec, ch);
     assert(err == VSL_STATUS_OK);
@@ -2306,14 +2326,18 @@ void irk_multinormal_vec_BM1(irk_state *state,
     if (len < 1)
         return;
 
-    while (len > 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_BOXMULLER, state->stream,
-                            MKL_INT_MAX, res, dim, storage_mode, mean_vec, ch);
+                            max_len, res, dim, storage_mode, mean_vec, ch);
         assert(err == VSL_STATUS_OK);
 
-        res += MKL_INT_MAX * dim;
-        len -= MKL_INT_MAX;
+        res += max_len * dim;
+        len -= max_len;
     }
 
     err = vdRngGaussianMV(VSL_RNG_METHOD_GAUSSIANMV_BOXMULLER, state->stream,
@@ -2335,14 +2359,18 @@ void irk_multinormal_vec_BM2(irk_state *state,
     if (len < 1)
         return;
 
-    while (len > 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_BOXMULLER2, state->stream,
-                            MKL_INT_MAX, res, dim, storage_mode, mean_vec, ch);
+                            max_len, res, dim, storage_mode, mean_vec, ch);
         assert(err == VSL_STATUS_OK);
 
-        res += MKL_INT_MAX * dim;
-        len -= MKL_INT_MAX;
+        res += max_len * dim;
+        len -= max_len;
     }
 
     err = vdRngGaussianMV(VSL_RNG_METHOD_GAUSSIANMV_BOXMULLER2, state->stream,

mkl_random/tests/test_random.py

@@ -1119,6 +1119,45 @@ def test_randomdist_multinormal_cholesky(randomdist):
     np.testing.assert_allclose(actual, desired, atol=1e-10, rtol=1e-10)
 
 
+@pytest.mark.parametrize("method", ["ICDF", "BoxMuller", "BoxMuller2"])
+def test_multinormal_cholesky_matches_split_calls(method):
+    # Results must not depend on how a draw is split across calls; the C
+    # layer relies on this when it splits large requests into chunks.
+    mean = np.array([0.1, -0.2])
+    chol_mat = np.array([[1.0, 0.0], [-0.5, 1.0]])
+    n1, n2 = 17, 29
+
+    whole = rnd.MKLRandomState(123).multinormal_cholesky(
+        mean, chol_mat, size=n1 + n2, method=method
+    )
+
+    split_state = rnd.MKLRandomState(123)
+    first = split_state.multinormal_cholesky(
+        mean, chol_mat, size=n1, method=method
+    )
+    second = split_state.multinormal_cholesky(
+        mean, chol_mat, size=n2, method=method
+    )
+
+    np.testing.assert_allclose(whole[:n1], first, rtol=1e-12, atol=1e-12)
+    np.testing.assert_allclose(whole[n1:], second, rtol=1e-12, atol=1e-12)
+
+
+@pytest.mark.parametrize("method", ["ICDF", "BoxMuller", "BoxMuller2"])
+@pytest.mark.parametrize("size,shape", [(0, (0, 2)), ((3, 0), (3, 0, 2))])
+def test_multinormal_cholesky_empty_size(method, size, shape):
+    # An empty draw returns an empty array and leaves the stream untouched.
+    mean = np.array([0.1, -0.2])
+    chol_mat = np.array([[1.0, 0.0], [-0.5, 1.0]])
+    reference = rnd.MKLRandomState(123)
+    state = rnd.MKLRandomState(123)
+    out = state.multinormal_cholesky(mean, chol_mat, size=size, method=method)
+    assert out.shape == shape
+    np.testing.assert_array_equal(
+        state.random_sample(32), reference.random_sample(32)
+    )
+
+
 def test_randomdist_negative_binomial(randomdist):
     rnd.seed(randomdist.seed, brng=randomdist.brng)
     actual = rnd.negative_binomial(n=100, p=0.12345, size=(3, 2))