Skip to content

Add oneMKL power loops for float32 and float64 - #285

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

antonwolfy wants to merge 9 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.
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