diff --git a/CHANGELOG.md b/CHANGELOG.md index 1e2fdd3e..bd96a346 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/mkl_umath/generate_umath.py b/mkl_umath/generate_umath.py index ca6a578b..15a80886 100644 --- a/mkl_umath/generate_umath.py +++ b/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, diff --git a/mkl_umath/src/mkl_umath_loops.c.src b/mkl_umath/src/mkl_umath_loops.c.src index c1a41660..398825fa 100644 --- a/mkl_umath/src/mkl_umath_loops.c.src +++ b/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# diff --git a/mkl_umath/src/mkl_umath_loops.h.src b/mkl_umath/src/mkl_umath_loops.h.src index 087a97cb..91d9342e 100644 --- a/mkl_umath/src/mkl_umath_loops.h.src +++ b/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 diff --git a/mkl_umath/src/ufuncsmodule.c b/mkl_umath/src/ufuncsmodule.c index 12564ab9..7625f355 100644 --- a/mkl_umath/src/ufuncsmodule.c +++ b/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; diff --git a/mkl_umath/tests/test_basic.py b/mkl_umath/tests/test_basic.py index 8433428c..e13a7e38 100644 --- a/mkl_umath/tests/test_basic.py +++ b/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}"