perf: batch array-valued distribution parameters into a single fill

CHANGELOG.md

@@ -12,6 +12,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 * Added a thread-safety section to the how-to guide for free-threaded Python [gh-159](https://github.com/IntelPython/mkl_random/pull/159)
 
 ### Changed
+* Sped up `normal`, `uniform`, `exponential`, `laplace`, `gumbel`, `logistic`, `rayleigh` and `lognormal` for array-valued parameters. Seeded results change for these array paths; scalar paths are unchanged [gh-171](https://github.com/IntelPython/mkl_random/pull/171)
+* Array parameters for these distributions must broadcast to the requested `size` without adding dimensions; previously accepted mismatches now raise `ValueError` [gh-171](https://github.com/IntelPython/mkl_random/pull/171)
+* `uniform` with array-valued bounds may return `high` due to floating-point rounding [gh-171](https://github.com/IntelPython/mkl_random/pull/171)
 * 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)
 * Sped up `randint` for power-of-two ranges at or above `INT_MAX` [gh-172](https://github.com/IntelPython/mkl_random/pull/172)
 * The random streams for power-of-two `randint` ranges at or above `INT_MAX` have changed: with a fixed seed these now produce different (but equally valid) samples. All other streams are unaffected. [gh-172](https://github.com/IntelPython/mkl_random/pull/172)

mkl_random/mklrand.pyx

@@ -757,6 +757,150 @@ cdef object vec_cont2_array(
     return arr_obj
 
 
+cdef object _param_out_shape(object size, tuple param_shapes):
+    """Result shape for a parameterised draw, matching the per-element paths."""
+    cdef object out_shape
+    cdef object bshape
+
+    if size is None:
+        return np.broadcast_shapes(*param_shapes)
+
+    out_shape = tuple(size) if np.iterable(size) else (size,)
+    try:
+        bshape = np.broadcast_shapes(out_shape, *param_shapes)
+    except ValueError:
+        raise ValueError("size is not compatible with inputs")
+    if bshape != out_shape:
+        raise ValueError("size is not compatible with inputs")
+    return out_shape
+
+
+cdef object _fill_standard2(
+    irk_state *state,
+    irk_cont2_vec func,
+    object out_shape,
+    object lock
+):
+    """Fill an entire request with one call, using standard parameters."""
+    cdef cnp.ndarray array
+    cdef cnp.npy_intp n
+    cdef double *array_data
+
+    array = <cnp.ndarray>np.empty(out_shape, np.float64)
+    n = cnp.PyArray_SIZE(array)
+    if n:
+        array_data = <double *>cnp.PyArray_DATA(array)
+        with lock, nogil:
+            func(state, n, array_data, 0.0, 1.0)
+    return array
+
+
+cdef object vec_loc_scale_array(
+    irk_state *state,
+    irk_cont2_vec func,
+    object size,
+    cnp.ndarray oloc,
+    cnp.ndarray oscale,
+    object lock
+):
+    """Draw a location and scale family with array-valued parameters.
+
+    ``func(0.0, 1.0)`` yields the standardised member, so
+    ``loc + scale * standardised`` is exact and needs one call per request.
+    """
+    cdef object array
+
+    array = _fill_standard2(
+        state,
+        func,
+        _param_out_shape(
+            size, ((<object>oloc).shape, (<object>oscale).shape)
+        ),
+        lock
+    )
+    np.multiply(array, oscale, out=array)
+    np.add(array, oloc, out=array)
+    return array
+
+
+cdef object vec_scale_array(
+    irk_state *state,
+    irk_cont1_vec func,
+    object size,
+    cnp.ndarray oscale,
+    object lock
+):
+    """Draw a scale family with an array-valued scale, one call per request."""
+    cdef cnp.ndarray array
+    cdef cnp.npy_intp n
+    cdef double *array_data
+
+    array = <cnp.ndarray>np.empty(
+        _param_out_shape(size, ((<object>oscale).shape,)), np.float64
+    )
+    n = cnp.PyArray_SIZE(array)
+    if n:
+        array_data = <double *>cnp.PyArray_DATA(array)
+        with lock, nogil:
+            func(state, n, array_data, 1.0)
+    np.multiply(array, oscale, out=array)
+    return array
+
+
+cdef object vec_uniform_array(
+    irk_state *state,
+    irk_cont2_vec func,
+    object size,
+    cnp.ndarray olow,
+    cnp.ndarray ohigh,
+    object lock
+):
+    """Draw uniforms over array-valued bounds, one call per request."""
+    cdef object array
+
+    array = _fill_standard2(
+        state,
+        func,
+        _param_out_shape(
+            size, ((<object>olow).shape, (<object>ohigh).shape)
+        ),
+        lock
+    )
+    np.multiply(array, np.subtract(ohigh, olow), out=array)
+    np.add(array, olow, out=array)
+    return array
+
+
+cdef object vec_lognormal_array(
+    irk_state *state,
+    irk_cont2_vec normal_func,
+    object size,
+    cnp.ndarray omean,
+    cnp.ndarray osigma,
+    object lock
+):
+    """Draw lognormals with array-valued parameters, one call per request.
+
+    Uses the normal fill: the parameters sit inside the exponential, so no
+    affine step applies to a standardised lognormal,
+    but exp(mean + sigma * z) does.
+    """
+    cdef object array
+
+    array = _fill_standard2(
+        state,
+        normal_func,
+        _param_out_shape(
+            size, ((<object>omean).shape, (<object>osigma).shape)
+        ),
+        lock
+    )
+    np.multiply(array, osigma, out=array)
+    np.add(array, omean, out=array)
+    np.exp(array, out=array)
+    return array
+
+
 cdef object vec_cont3_array_sc(
     irk_state *state,
     irk_cont3_vec func,
@@ -2730,16 +2874,19 @@ cdef class _MKLRandomState:
         Samples are uniformly distributed over the half-open interval
         ``[low, high)`` (includes low, but excludes high).  In other words,
         any value within the given interval is equally likely to be drawn
-        by `uniform`.
+        by `uniform`. With array-valued bounds, floating-point rounding
+        may include the upper boundary in the returned samples.
 
         Parameters
         ----------
         low : float, optional
             Lower boundary of the output interval.  All values generated will be
             greater than or equal to low.  The default value is 0.
         high : float
-            Upper boundary of the output interval.  All values generated will be
-            less than high.  The default value is 1.0.
+            Upper boundary of the output interval. With array-valued bounds,
+            high may be included due to floating-point rounding in
+            ``low + (high - low) * U``, where ``U`` is drawn from ``[0, 1)``.
+            The default value is 1.0.
         size : int or tuple of ints, optional
             Output shape.  If the given shape is, e.g., ``(m, n, k)``, then
             ``m * n * k`` samples are drawn.  Default is None, in which case a
@@ -2824,7 +2971,7 @@ cdef class _MKLRandomState:
         if np.any(olow >= ohigh):
             raise ValueError("low >= high")
 
-        return vec_cont2_array(
+        return vec_uniform_array(
             self.internal_state, irk_uniform_vec, size, olow, ohigh, self.lock
         )
 
@@ -3226,23 +3373,23 @@ cdef class _MKLRandomState:
             method, [ICDF, BOXMULLER, BOXMULLER2], _method_alias_dict_gaussian
         )
         if method is ICDF:
-            return vec_cont2_array(
+            return vec_loc_scale_array(
                 self.internal_state,
                 irk_normal_vec_ICDF,
                 size,
                 oloc,
                 oscale, self.lock
             )
         elif method is BOXMULLER2:
-            return vec_cont2_array(
+            return vec_loc_scale_array(
                 self.internal_state,
                 irk_normal_vec_BM2,
                 size,
                 oloc,
                 oscale, self.lock
             )
         else:
-            return vec_cont2_array(
+            return vec_loc_scale_array(
                 self.internal_state,
                 irk_normal_vec_BM1,
                 size,
@@ -3381,7 +3528,7 @@ cdef class _MKLRandomState:
 
         if np.any(np.signbit(oscale) | (oscale == 0)):
             raise ValueError("scale <= 0")
-        return vec_cont1_array(
+        return vec_scale_array(
             self.internal_state, irk_exponential_vec, size, oscale, self.lock
         )
 
@@ -4826,8 +4973,9 @@ cdef class _MKLRandomState:
 
         if np.any(np.signbit(oscale) | np.equal(oscale, 0.0)):
             raise ValueError("scale <= 0")
-        return vec_cont2_array(
-            self.internal_state, irk_laplace_vec, size, oloc, oscale, self.lock
+        return vec_loc_scale_array(
+            self.internal_state, irk_laplace_vec, size, oloc, oscale,
+            self.lock
         )
 
     def gumbel(self, loc=0.0, scale=1.0, size=None):
@@ -4966,8 +5114,9 @@ cdef class _MKLRandomState:
 
         if np.any(np.signbit(oscale) | np.equal(oscale, 0.0)):
             raise ValueError("scale <= 0")
-        return vec_cont2_array(
-            self.internal_state, irk_gumbel_vec, size, oloc, oscale, self.lock
+        return vec_loc_scale_array(
+            self.internal_state, irk_gumbel_vec, size, oloc, oscale,
+            self.lock
         )
 
     def logistic(self, loc=0.0, scale=1.0, size=None):
@@ -5067,7 +5216,7 @@ cdef class _MKLRandomState:
 
         if np.any(np.signbit(oscale) | np.equal(oscale, 0.0)):
             raise ValueError("scale <= 0")
-        return vec_cont2_array(
+        return vec_loc_scale_array(
             self.internal_state,
             irk_logistic_vec,
             size,
@@ -5228,18 +5377,18 @@ cdef class _MKLRandomState:
             method, [ICDF, BOXMULLER], _method_alias_dict_gaussian_short
         )
         if method is ICDF:
-            return vec_cont2_array(
+            return vec_lognormal_array(
                 self.internal_state,
-                irk_lognormal_vec_ICDF,
+                irk_normal_vec_ICDF,
                 size,
                 omean,
                 osigma,
                 self.lock
             )
         else:
-            return vec_cont2_array(
+            return vec_lognormal_array(
                 self.internal_state,
-                irk_lognormal_vec_BM,
+                irk_normal_vec_BM2,
                 size,
                 omean,
                 osigma,
@@ -5321,7 +5470,7 @@ cdef class _MKLRandomState:
 
         if np.any(np.signbit(oscale) | np.equal(oscale, 0.0)):
             raise ValueError("scale <= 0.0")
-        return vec_cont1_array(
+        return vec_scale_array(
             self.internal_state, irk_rayleigh_vec, size, oscale, self.lock
         )