Skip to content

Add oneMKL power loops for float32 and float64 - #285

Open
antonwolfy wants to merge 10 commits into
mainfrom
add-power-ufunc
Open

antonwolfy wants to merge 10 commits into
mainfrom
add-power-ufunc

Conversation

@antonwolfy

@antonwolfy antonwolfy commented Oct 14, 2026 •

Copy link
Copy Markdown
Collaborator

This PR adds power to mkl_umath for float32 and float64. Contiguous arrays above the usual 8192-element threshold go to v?Pow, a scalar exponent is broadcast with v?PowI and a zero increment, and everything else (small, strided, scalar base) stays on the pow loop.

Matching NumPy semantics

v?Pow returns the same values as pow on all IEEE special inputs (signed zeros, infinities, NaN, subnormals, negative bases) and is within 1 ULP of libm on random inputs. A few MKL behaviours needed handling to make the loop a drop-in replacement:

  • Scalar exponent fast path. Since 2.3, NumPy computes x ** -1, 0, 0.5, 1, 2 directly (0.5 uses sqrt semantics, so (-0.0) ** 0.5 is -0.0, not +0.0 as with pow). The loop mirrors this when built against NumPy >= 2.3.
  • v?Powx is not used. For exponents -0.5, 0.5 and 1.5 it returns wrong values for -0.0 and -inf (e.g. (-inf) ** 1.5 gives nan). v?PowI with a zero exponent increment uses the same kernel as v?Pow, so its values are correct, and it is as fast as v?Powx for other exponents.
  • Floating-point flags. MKL's flags cannot be relied on: some are spurious (0 ** -1 also raises invalid), flags raised in MKL worker threads never reach the calling thread, and on some CPUs (seen on CI runners, most likely AVX-512) float32 domain errors such as (-1.0) ** 1.5 raise nothing at all, not even in the VM error status. So after every MKL call the loop clears MKL's flags, scans the results for nan and inf (every flag-worthy case produces one), and only then derives the flags C99 requires, as glibc raises them, from inputs and results. This makes the warnings the same on every CPU and thread count. 0 ** -inf raises nothing, since C99 allows but does not require divide by zero there.
  • In-place calls (out aliasing an input) go through a 64K-element buffer so the inputs are still available for that.

Complex loops are not added: v?Pow computes exp(y*log(x)), which is not exact for integer exponents (1j**2 gives -1+1.2e-16j instead of -1+0j).

Performance

Measured on a Lunar Lake DUT (4 P-cores + 4 LP E-cores, AVX2, no AVX-512) with NumPy 2.5.4 and oneMKL 2026.2. Stock and patched np.power were timed in the same process on the same arrays, with the performance governor, turbo off, single-thread runs pinned to a P-core and 4-thread runs to the P-cores. An A/A control (stock vs stock across a patch/restore cycle) read 1.00x in every cell.

Speedup of patched over stock np.power at 1M elements:

Case 1 thread 4 threads 8 threads
float32, array ** array 2.6x 8.7x 9.0x
float32, x ** scalar (2.5, 1.5, -0.5) 2.6x 8.6x 9.2x
float64, array ** array 2.6x 8.2x 7.8x
float64, x ** scalar (2.5, 1.5, -0.5) 2.2x 7.7x 7.7x

At 10k elements, just above the threshold, the speedup is 2.3-2.6x single-threaded and about 5x threaded; at 10M it reaches 11x for float32 at 8 threads. The results scan runs on the calling thread: it costs 4-11% single-threaded and 30-45% at 4-8 threads for 1M+ elements, where the speedups would otherwise be 10-15x. Not measured yet on AVX-512 Xeons, where NumPy's own power uses SVML, so the gap there will be smaller.

Testing

  • New tests in test_basic.py: special values for array and scalar exponents on both the fallback and the MKL path, in-place on either operand, and the warnings C99 requires (spelled out rather than compared with np.power, whose warnings for these edge cases differ between CPUs), including a bad element at the end of a large array, processed by an MKL worker thread.
  • pytest --pyargs mkl_umath: 432 passed, 1 skipped, locally and on the DUT, at 1, 4 and 8 MKL threads.
  • NumPy's test_umath.py with NumPy patched: the only failures are 12 test_unary_spurious_fpexception cases for sign, which fail identically on main and are unrelated to this change.

jharlow-intel
jharlow-intel previously approved these changes Oct 15, 2026

@jharlow-intel jharlow-intel left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks, this is really nice work! The numerics and FP-flag handling look right, and the testing is thorough: special values, in-place calls, and flags raised from MKL worker threads are all covered. A couple of small things inline.

Comment thread mkl_umath/src/mkl_umath_loops.c.src Outdated
Comment thread CHANGELOG.md
Comment thread mkl_umath/tests/test_basic.py Outdated
Arrays use v?Pow and a scalar exponent uses v?Powx, above the usual
8192-element threshold. The loop keeps NumPy's semantics:

- the scalar exponent fast path of NumPy >= 2.3 (-1, 0, 0.5, 1, 2) is
  mirrored when built against it
- v?Powx gets -0.0 and -inf wrong for exponents -0.5, 0.5 and 1.5, so
  those elements are recomputed with pow
- MKL raises spurious invalid/overflow flags for some special inputs and
  drops flags raised in its worker threads, so when the VM error status
  or the flags report a problem they are recomputed with pow
- in-place calls go through a buffer so the inputs survive for that

Complex loops are left to NumPy: v?Pow is not exact for integer
exponents (1j**2).
test_power_warnings failed on some CI runners: it compared against
np.power, whose warnings for edge cases such as 0 ** -inf depend on
NumPy's CPU dispatch, and the flag recheck re-ran pow, whose flags depend
on the math library and CPU as well.

- derive the flags in the recheck from inputs and results as C99 (and
  glibc) raise them, instead of calling pow again
- broadcast a scalar exponent with v?PowI and a zero increment instead of
  v?Powx, which never sets the VM error status (so flags from worker
  threads were lost) and needed a fix-up for -0.0 and -inf
- spell out the expected warnings in the test, and report them on failure
On some CI runners (most likely AVX-512) MKL's float32 v?Pow raises
neither a flag nor the VM error status for a domain error such as
(-1.0) ** 1.5, so the warning was lost. Instead of trusting MKL to report
problems, clear its flags after every call, scan the results for nan and
inf, and only then derive the C99 flags from inputs and results.
On the Windows Jenkins jobs NumPy returns 0 for float32 1e-30 ** 1.5 and
1e30 ** -1.5, where MKL returns the correctly rounded smallest subnormal.
A relative tolerance cannot absorb that, so add the dtype's smallest
subnormal as absolute tolerance.
The fast path mirrors NumPy >= 2.3, but was gated on the NumPy headers
mkl_umath was built against. A build against 2.3+ headers running on
NumPy <= 2.2 then gave sqrt semantics for an exponent of 0.5, e.g.
(-inf)**0.5 = nan, and a build against older headers running on 2.3+
did not. _ufuncs now passes the runtime C-API feature version to the
loops library on import.
@antonwolfy

Copy link
Copy Markdown
Collaborator Author

I reran the performance measurements on a second Lunar Lake DUT, and the gains in the description reproduce.

Setup: Intel Core Ultra 5 236V (4 P-cores + 4 LP E-cores, AVX2, no AVX-512), Ubuntu 25.10, PR head d3a7cd3 built with icx against NumPy 2.5.4 and oneMKL 2026.2 (Python 3.13). Same method as in the description: stock and patched np.power timed in the same process on the same arrays (best of 7), performance governor, turbo off, 1 thread pinned to P-core 0, 4 threads to P-cores 0-3, 8 threads on all cores. Bases are uniform in [0.1, 10], array exponents uniform in [0.5, 3], and patched results are checked against stock with allclose. The A/A control (stock vs stock across a patch/restore cycle) was run at 1, 4 and 8 threads and read 1.00-1.01x in every cell. pytest --pyargs mkl_umath: 432 passed, 1 skipped at 1, 4 and 8 MKL threads.

Speedup of patched over stock np.power at 1M elements, description → this run (range over two runs; the scalar rows cover exponents 2.5, 1.5 and -0.5):

Case 1 thread 4 threads 8 threads
float32, array ** array 2.6x → 2.61-2.62x 8.7x → 8.54x 9.0x → 9.98-10.03x
float32, x ** scalar 2.6x → 2.52-2.55x 8.6x → 8.07-8.14x 9.2x → 9.23-9.92x
float64, array ** array 2.6x → 2.58-2.59x 8.2x → 7.94-7.99x 7.8x → 7.56-7.57x
float64, x ** scalar 2.2x → 2.36-2.38x 7.7x → 7.39-7.42x 7.7x → 7.33-7.43x

Every cell is within -6% to +11% of the description.

Other claims:

  • 10k elements, single-threaded: 2.32-2.57x (description: 2.3-2.6x).
  • 10k elements, threaded: 4.6-5.1x at 4 threads and 4.5-4.9x for float32 at 8 threads, in line with "about 5x". The exception is float64 at 8 threads, at 2.4-2.9x; the original DUT run showed the same (2.9x for array, 2.5x for scalar exponents), so "about 5x threaded" in the description does not cover that cell.
  • 10M elements, float32 at 8 threads: 10.7-11.0x (description: 11x).
  • Cost of the results scan, measured against a build with the isfinite loop removed (16 warning tests fail with it, so the scan was really gone), alternating the two builds in separate processes; runs of the same build differed by at most 5%. The scan costs 5-10% single-threaded and 27-51% at 4-8 threads for 1M+ elements, where the build without it reaches 10-15x (description: 4-11%, 30-45%, 10-15x). At 10k-100k elements with threads it costs more, 41-145%, and that is what holds back float64 at 8 threads and 10k: without the scan that cell is 6.0-6.2x.
Full tables, run 1: speedup (patched time)

1 thread(s) on CPUs 0

dtype    exponent           10,000          100,000        1,000,000       10,000,000
float32     array   2.56x (  0.03ms)   2.73x (  0.31ms)   2.62x (  3.26ms)   2.65x ( 32.17ms)
float32       2.5   2.51x (  0.03ms)   2.59x (  0.33ms)   2.54x (  3.35ms)   2.47x ( 34.49ms)
float32       3.0   2.50x (  0.03ms)   2.59x (  0.33ms)   2.54x (  3.36ms)   2.47x ( 34.49ms)
float32       1.5   2.50x (  0.03ms)   2.59x (  0.33ms)   2.53x (  3.36ms)   2.47x ( 34.49ms)
float32      -0.5   2.51x (  0.03ms)   2.59x (  0.33ms)   2.54x (  3.35ms)   2.46x ( 34.60ms)
float64     array   2.56x (  0.05ms)   2.66x (  0.49ms)   2.58x (  5.01ms)   2.58x ( 50.00ms)
float64       2.5   2.32x (  0.05ms)   2.39x (  0.54ms)   2.37x (  5.46ms)   2.31x ( 55.90ms)
float64       3.0   2.32x (  0.05ms)   2.41x (  0.54ms)   2.37x (  5.46ms)   2.30x ( 55.97ms)
float64       1.5   2.32x (  0.05ms)   2.39x (  0.54ms)   2.37x (  5.46ms)   2.31x ( 55.79ms)
float64      -0.5   2.32x (  0.05ms)   2.40x (  0.54ms)   2.36x (  5.46ms)   2.30x ( 55.97ms)

4 thread(s) on CPUs 0-3

dtype    exponent           10,000          100,000        1,000,000       10,000,000
float32     array   4.81x (  0.02ms)   6.76x (  0.13ms)   8.54x (  1.00ms)   8.72x (  9.79ms)
float32       2.5   4.56x (  0.02ms)   5.97x (  0.14ms)   8.09x (  1.05ms)   8.30x ( 10.27ms)
float32       3.0   4.56x (  0.02ms)   5.96x (  0.14ms)   8.14x (  1.05ms)   8.32x ( 10.25ms)
float32       1.5   4.58x (  0.02ms)   6.07x (  0.14ms)   8.14x (  1.05ms)   8.30x ( 10.27ms)
float32      -0.5   4.58x (  0.02ms)   6.14x (  0.14ms)   8.14x (  1.05ms)   8.27x ( 10.32ms)
float64     array   5.13x (  0.02ms)   7.06x (  0.18ms)   7.94x (  1.63ms)   8.12x ( 15.92ms)
float64       2.5   4.79x (  0.03ms)   6.17x (  0.21ms)   7.41x (  1.74ms)   7.49x ( 17.21ms)
float64       3.0   4.77x (  0.03ms)   6.11x (  0.21ms)   7.39x (  1.75ms)   7.49x ( 17.22ms)
float64       1.5   4.79x (  0.03ms)   6.19x (  0.21ms)   7.42x (  1.74ms)   7.51x ( 17.18ms)
float64      -0.5   4.80x (  0.03ms)   6.15x (  0.21ms)   7.41x (  1.74ms)   7.51x ( 17.16ms)

8 thread(s) on CPUs 0-7

dtype    exponent           10,000          100,000        1,000,000       10,000,000
float32     array   4.85x (  0.02ms)   7.66x (  0.11ms)   9.98x (  0.85ms)  10.92x (  8.11ms)
float32       2.5   4.53x (  0.02ms)   7.58x (  0.11ms)   9.73x (  0.88ms)  10.75x (  8.02ms)
float32       3.0   4.54x (  0.02ms)   7.53x (  0.11ms)   9.40x (  0.90ms)  10.68x (  8.04ms)
float32       1.5   4.55x (  0.02ms)   7.56x (  0.11ms)   9.23x (  0.92ms)  10.73x (  8.00ms)
float32      -0.5   4.57x (  0.02ms)   7.57x (  0.11ms)   9.34x (  0.91ms)  10.77x (  8.00ms)
float64     array   2.87x (  0.04ms)   6.17x (  0.21ms)   7.56x (  1.71ms)   7.71x ( 16.76ms)
float64       2.5   2.45x (  0.05ms)   5.96x (  0.22ms)   7.38x (  1.75ms)   7.71x ( 16.76ms)
float64       3.0   2.45x (  0.05ms)   5.88x (  0.22ms)   7.44x (  1.74ms)   7.75x ( 16.64ms)
float64       1.5   2.42x (  0.05ms)   5.91x (  0.22ms)   7.36x (  1.75ms)   7.76x ( 16.62ms)
float64      -0.5   2.46x (  0.05ms)   5.88x (  0.22ms)   7.43x (  1.74ms)   7.78x ( 16.58ms)

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.

2 participants