Add oneMKL power loops for float32 and float64

CHANGELOG.md

@@ -4,6 +4,11 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [dev] - YYYY-MM-DD
+
+### Added
+* Added oneMKL-backed `power` loops for float32 and float64 arrays, using `v?Pow`, and `v?PowI` to broadcast a scalar exponent [gh-285](https://github.com/IntelPython/mkl_umath/pull/285)
+
 ## [0.6.0] - 2026-10-15
 
 ### Added

mkl_umath/generate_umath.py

@@ -501,7 +501,16 @@ def english_upper(s):
         TD(inexactvec + cmplxvec),
     ),
     # '_ones_like':
-    # 'power':
+    # 'power': complex loops are not added since v?Pow computes exp(y*log(x)),
+    #          which is not exact for integer exponents (e.g. 1j**2)
+    "power": Ufunc(
+        2,
+        1,
+        None,
+        docstrings.get("numpy._core.umath.power"),
+        None,
+        TD(inexactvec),
+    ),
     # 'float_power':
     "absolute": Ufunc(
         1,

mkl_umath/src/mkl_umath_loops.c.src

@@ -71,6 +71,25 @@
 #define VML_ASM_THRESHOLD 100000
 #define VML_D_THRESHOLD 8000
 
+/* floating-point exceptions reported by NumPy that v?Pow does not raise reliably */
+#define POW_FPE_FLAGS (FE_INVALID | FE_DIVBYZERO | FE_OVERFLOW)
+/* block size, in elements, for in-place power, which goes through a buffer */
+#define POW_INPLACE_BLOCK ((npy_intp) 1 << 16)
+
+/* not defined by the NumPy < 2.3 headers */
+#ifndef NPY_2_3_API_VERSION
+#define NPY_2_3_API_VERSION 0x00000014
+#endif
+
+/* C-API feature version of the NumPy loaded at runtime, set by _ufuncs on import */
+static int numpy_runtime_version = 0;
+
+void
+mkl_umath_set_numpy_runtime_version(int version)
+{
+    numpy_runtime_version = version;
+}
+
 #define MKL_INT_MAX ((npy_intp) ((~((MKL_UINT) 0)) >> 1))
 
 #define CHUNKED_VML_CALL2(vml_func, n, type, mkl_type, mkl_ftype, in1, op1)     \
@@ -107,6 +126,23 @@
         }                                                                       \
     } while(0)
 
+/* strided VM call with a zero increment for in2, which broadcasts its first element */
+#define CHUNKED_VML_CALL3_S2(vml_func, n, type, in1, in2, op1)                 \
+    do  {                                                                       \
+        npy_intp _n_ = (n);                                                     \
+        const npy_intp _chunk_size = MKL_INT_MAX;                               \
+        type *in1p = (type *) (in1);                                            \
+        type *in2p = (type *) (in2);                                            \
+        type *op1p = (type *) (op1);                                            \
+        while (_n_ > 0) {                                                       \
+            npy_intp _current_chunk = (_n_ > _chunk_size) ? _chunk_size : _n_;  \
+            vml_func((MKL_INT) _current_chunk, in1p, 1, in2p, 0, op1p, 1);      \
+            _n_ -= _current_chunk;                                              \
+            in1p += _current_chunk;                                             \
+            op1p += _current_chunk;                                             \
+        }                                                                       \
+    } while(0)
+
 
 #define CHUNKED_VML_LINEARFRAC_CALL(vml_func, n, type, in1, op1, scaleA, shiftA, scaleB, shiftB) \
      do {                                                                                        \
@@ -837,6 +873,145 @@ mkl_umath_@TYPE@_@kind@(char **args, const npy_intp *dimensions, const npy_intp
 }
 /**end repeat1**/
 
+/*
+ * v@s@Pow returns the same values as pow@c@, a scalar exponent is broadcast
+ * with v@s@PowI rather than passed to v@s@Powx, which gets -0.0 and -inf
+ * wrong for some exponents. The floating-point exception flags MKL raises
+ * cannot be relied on though: they are spurious for some special inputs
+ * (e.g. 0 ** -1 also raises invalid), those raised in MKL worker threads never
+ * reach the calling thread, and on some CPUs float32 domain errors raise
+ * nothing, not even in the VM error status. So MKL's flags are replaced with
+ * the ones C99 (and glibc) raise, derived from the inputs and results: each
+ * such case gives a nan or an inf, which a cheap scan of the results finds.
+ * r must not alias x or y, which are read again.
+ */
+static void
+@TYPE@_power_vml(const @type@ *x, const @type@ *y, const int y_is_scalar, @type@ *r, const npy_intp n)
+{
+    /* flags raised by earlier calls of the loop are NumPy's to report */
+    const int prev_flags = fetestexcept(POW_FPE_FLAGS);
+    int nonfinite = 0;
+    npy_intp i;
+
+    if (y_is_scalar) {
+        CHUNKED_VML_CALL3_S2(v@s@PowI, n, @type@, x, y, r);
+    }
+    else {
+        CHUNKED_VML_CALL3(v@s@Pow, n, @type@, @type@, x, y, r);
+    }
+    feclearexcept(POW_FPE_FLAGS);
+
+    /* branch-free so that it vectorizes */
+    for (i = 0; i < n; i++) {
+        nonfinite |= !isfinite(r[i]);
+    }
+    if (nonfinite) {
+        int flags = 0;
+
+        for (i = 0; i < n; i++) {
+            const @type@ in1 = x[i];
+            const @type@ in2 = y_is_scalar ? *y : y[i];
+
+            if (isnan(r[i]) && !isnan(in1) && !isnan(in2)) {
+                flags |= FE_INVALID;
+            }
+            else if (isinf(r[i]) && isfinite(in1) && isfinite(in2)) {
+                /* 0 ** -inf is +inf too, but C99 does not require a flag there */
+                flags |= (in1 == 0) ? FE_DIVBYZERO : FE_OVERFLOW;
+            }
+        }
+        if (flags) {
+            feraiseexcept(flags);
+        }
+    }
+    if (prev_flags) {
+        feraiseexcept(prev_flags);
+    }
+}
+
+void
+mkl_umath_@TYPE@_power(char **args, const npy_intp *dimensions, const npy_intp *steps, void *NPY_UNUSED(func))
+{
+    const npy_intp n = dimensions[0];
+
+    /*
+     * same fast path for a scalar exponent as in NumPy >= 2.3, checked against
+     * the NumPy loaded rather than the headers mkl_umath was built against
+     */
+    if (steps[1] == 0 && numpy_runtime_version >= NPY_2_3_API_VERSION) {
+        const @type@ in2 = *(@type@ *)args[1];
+
+        if (in2 == -1.0 || in2 == 0.0 || in2 == 0.5 || in2 == 1.0 || in2 == 2.0) {
+            BINARY_LOOP {
+                const @type@ in1 = *(@type@ *)ip1;
+
+                if (in2 == -1.0) {
+                    *(@type@ *)op1 = 1.0 / in1;
+                }
+                else if (in2 == 0.0) {
+                    *(@type@ *)op1 = 1.0;
+                }
+                else if (in2 == 0.5) {
+                    *(@type@ *)op1 = sqrt@c@(in1);
+                }
+                else if (in2 == 1.0) {
+                    *(@type@ *)op1 = in1;
+                }
+                else {
+                    *(@type@ *)op1 = in1 * in1;
+                }
+            }
+            return;
+        }
+    }
+
+    {
+        const int contig = IS_BINARY_CONT(@type@, @type@);
+        const int contig_scalar = IS_BINARY_CONT_S2(@type@, @type@);
+        const int disjoint_or_same1 = DISJOINT_OR_SAME(args[0], args[2], n, sizeof(@type@));
+        const int disjoint_or_same2 = DISJOINT_OR_SAME(args[1], args[2], n, sizeof(@type@));
+
+        if (n > VML_TRANSCEDENTAL_THRESHOLD && disjoint_or_same1 &&
+            ((contig && disjoint_or_same2) || contig_scalar)) {
+            const @type@ *x = (const @type@ *)args[0];
+            const @type@ *y = (const @type@ *)args[1];
+            @type@ *r = (@type@ *)args[2];
+            const int y_is_scalar = !contig;
+
+            if (r != x && (y_is_scalar || r != y)) {
+                @TYPE@_power_vml(x, y, y_is_scalar, r, n);
+                return;
+            }
+            else {
+                /* in-place: keep the inputs intact until their block is done */
+                const npy_intp block = (n < POW_INPLACE_BLOCK) ? n : POW_INPLACE_BLOCK;
+                @type@ *buf = (@type@ *) malloc(block * sizeof(@type@));
+
+                if (buf != NULL) {
+                    npy_intp off;
+
+                    for (off = 0; off < n; off += block) {
+                        const npy_intp len = (n - off < block) ? n - off : block;
+
+                        @TYPE@_power_vml(x + off, y_is_scalar ? y : y + off, y_is_scalar, buf, len);
+                        memcpy(r + off, buf, len * sizeof(@type@));
+                    }
+                    free(buf);
+                    return;
+                }
+            }
+        }
+    }
+
+    {
+        BINARY_LOOP {
+            const @type@ in1 = *(@type@ *)ip1;
+            const @type@ in2 = *(@type@ *)ip2;
+            *((@type@ *)op1) = pow@c@(in1, in2);
+        }
+    }
+}
+
 /**begin repeat1
  * # kind = cos, sin, tan, arccos, arcsin, arctan, cosh, sinh, tanh, arccosh, arcsinh, arctanh, fabs, floor, ceil, rint, trunc, cbrt, sqrt, expm1, log, log2, log1p, log10#
  * # func = cos, sin, tan, acos, asin, atan, cosh, sinh, tanh, acosh, asinh, atanh, fabs, floor, ceil, rint, trunc, cbrt, sqrt, expm1, log, log2, log1p, log10#

mkl_umath/src/mkl_umath_loops.h.src

@@ -48,6 +48,14 @@
     #define USE_NUMPY_2
 #endif
 
+/*
+ * Records the C-API feature version of the NumPy loaded at runtime, for the
+ * loops whose results depend on it (power in NumPy >= 2.3).
+ */
+MKL_UMATH_API
+void
+mkl_umath_set_numpy_runtime_version(int version);
+
 /**begin repeat
  * Float types
  *  #TYPE = FLOAT, DOUBLE#
@@ -58,7 +66,7 @@
             copysign, cos, cosh, divide, divmod, isfinite, isinf, isnan, equal, exp, exp2, expm1, fabs,
             floor, fmax, fmin, frexp, greater, greater_equal, ldexp, less, less_equal, log, log2, log10,
             log1p, logical_and, logical_not, logical_or, logical_xor, multiply, modf,
-            negative, nextafter, not_equal, positive, reciprocal, rint, sign, signbit, sin,
+            negative, nextafter, not_equal, positive, power, reciprocal, rint, sign, signbit, sin,
             sinh, spacing, sqrt, square, subtract, tan, tanh, trunc#
  */
 MKL_UMATH_API

mkl_umath/src/ufuncsmodule.c

@@ -45,6 +45,10 @@ PyMODINIT_FUNC PyInit__ufuncs(void)
     import_array();
     import_umath();
 
+    /* PyArray_RUNTIME_VERSION, which the NumPy 1.x headers do not provide */
+    mkl_umath_set_numpy_runtime_version(
+        (int)PyArray_GetNDArrayCFeatureVersion());
+
     m = PyModule_Create(&_ufuncs_module);
     if (m == NULL)
         return NULL;

mkl_umath/tests/test_basic.py

@@ -23,6 +23,8 @@
 # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
+import warnings
+
 import numpy as np
 import pytest
 
@@ -70,7 +72,7 @@ def get_args(args_str, size, low, high):
         if umath in ["arccos", "arcsin", "arctanh"]:
             low = -1
             high = 1
-        elif umath in ["log", "log10", "log1p", "log2", "sqrt"]:
+        elif umath in ["log", "log10", "log1p", "log2", "power", "sqrt"]:
             low = 0
             high = 10
         elif umath == "arccosh":
@@ -202,3 +204,85 @@ def test_absolute_nan_signbit(size, dtype):
             f"absolute({sign}nan) must be +nan for "
             f"{dtype.__name__} at size {size}"
         )
+
+
+_power_specials = [0.0, -0.0, 0.5, -0.5, 1.0, -1.0, 1.5, -1.5, 2.0, -2.0]
+_power_specials += [3.0, -3.0, 1e-30, 1e30, np.inf, -np.inf, np.nan]
+
+
+def _assert_power_equal(mkl_res, np_res):
+    # MKL may differ from libm by 1 ULP, also for subnormal results (NumPy on
+    # Windows returns 0 for float32 1e-30 ** 1.5), special values must match
+    # exactly
+    one = np_res.dtype.type(1)
+    eps, tiny = np.spacing(one), np.nextafter(0 * one, one)
+    np.testing.assert_allclose(
+        mkl_res, np_res, rtol=2 * eps, atol=tiny, equal_nan=True
+    )
+    assert (np.signbit(mkl_res) == np.signbit(np_res))[~np.isnan(np_res)].all()
+
+
+def _power_warnings(func, *args):
+    with warnings.catch_warnings(record=True) as w:
+        warnings.simplefilter("always")
+        func(*args)
+    return sorted(str(x.message).split(" encountered")[0] for x in w)
+
+
+@pytest.mark.parametrize("size", [100, 8192 + 1])
+@pytest.mark.parametrize("dtype", [np.float32, np.float64])
+def test_power_special_values(size, dtype):
+    # size 100 takes the fallback loop, 8193 takes MKL's v?Pow and v?PowI
+    specials = np.array(_power_specials, dtype=dtype)
+    x, y = (a.ravel() for a in np.meshgrid(specials, specials))
+    with np.errstate(all="ignore"):
+        x = np.resize(x, max(size, x.size))
+        y = np.resize(y, max(size, y.size))
+        _assert_power_equal(mu.power(x, y), np.power(x, y))
+        for e in specials:
+            _assert_power_equal(mu.power(x, e), np.power(x, e))
+
+
+@pytest.mark.parametrize("dtype", [np.float32, np.float64])
+@pytest.mark.parametrize("exponent", [None, 2.5, 1.5, -0.5])
+def test_power_inplace(dtype, exponent):
+    x = np.random.uniform(0, 10, 10**5 + 1).astype(dtype)
+    y = np.random.uniform(-3, 3, 10**5 + 1).astype(dtype)
+    if exponent is not None:
+        y = dtype(exponent)
+    np_res = np.power(x, y)
+
+    mkl_res = x.copy()
+    mu.power(mkl_res, y, out=mkl_res)
+    _assert_power_equal(mkl_res, np_res)
+    if exponent is None:
+        mkl_res = y.copy()
+        mu.power(x, mkl_res, out=mkl_res)
+        _assert_power_equal(mkl_res, np_res)
+
+
+@pytest.mark.parametrize("dtype", [np.float32, np.float64])
+@pytest.mark.parametrize(
+    "base, exponent, expected",
+    [
+        (2.0, 1.5, []),
+        (-1.0, 1.5, ["invalid value"]),
+        (0.0, -1.0, ["divide by zero"]),
+        (-0.0, -0.5, ["divide by zero"]),
+        (1e30, 20.0, ["overflow"]),
+        # C99 allows, but does not require, divide by zero here
+        (0.0, -np.inf, []),
+    ],
+)
+@pytest.mark.parametrize("pos", [0, -1])
+def test_power_warnings(dtype, base, exponent, expected, pos):
+    # The flags C99 requires, as glibc raises them. They are spelled out since
+    # NumPy's own loops do not raise them the same way on every CPU. MKL raises
+    # spurious flags for some of these, and flags raised in its worker threads
+    # (for the element at the end) never reach the caller.
+    x = np.full(10**6, 2.0, dtype=dtype)
+    y = np.full(10**6, 1.5, dtype=dtype)
+    x[pos], y[pos] = base, exponent
+    for args in ((x, y), (x, y[pos])):
+        got = _power_warnings(mu.power, *args)
+        assert got == expected, f"power{(base, exponent)} warned {got}"