Repository navigation
Add oneMKL power loops for float32 and float64 - #285
Open
antonwolfy wants to merge 9 commits into
Open
antonwolfy wants to merge 9 commits into
antonwolfy wants to merge 9 commits into
Conversation
antonwolfy
requested review from
jharlow-intel,
ndgrigorian,
vlad-perevezentsev and
xaleryb
as code owners
October 14, 2026 14:32
jharlow-intel
previously approved these changes
Oct 15, 2026
jharlow-intel
left a comment
Collaborator
There was a problem hiding this comment.
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.
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
force-pushed
the
add-power-ufunc
branch
from
October 15, 2026 18:07
3f8c596 to
2582d50
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This PR adds
powertomkl_umathfor float32 and float64. Contiguous arrays above the usual 8192-element threshold go tov?Pow, a scalar exponent is broadcast withv?PowIand a zero increment, and everything else (small, strided, scalar base) stays on thepowloop.Matching NumPy semantics
v?Powreturns the same values aspowon 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:x ** -1, 0, 0.5, 1, 2directly (0.5 usessqrtsemantics, so(-0.0) ** 0.5is-0.0, not+0.0as withpow). The loop mirrors this when built against NumPy >= 2.3.v?Powxis not used. For exponents -0.5, 0.5 and 1.5 it returns wrong values for-0.0and-inf(e.g.(-inf) ** 1.5givesnan).v?PowIwith a zero exponent increment uses the same kernel asv?Pow, so its values are correct, and it is as fast asv?Powxfor other exponents.0 ** -1also 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.5raise 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 ** -infraises nothing, since C99 allows but does not require divide by zero there.outaliasing an input) go through a 64K-element buffer so the inputs are still available for that.Complex loops are not added:
v?Powcomputesexp(y*log(x)), which is not exact for integer exponents (1j**2gives-1+1.2e-16jinstead 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.powerwere timed in the same process on the same arrays, with theperformancegovernor, 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.powerat 1M elements: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
poweruses SVML, so the gap there will be smaller.Testing
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 withnp.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.test_umath.pywith NumPy patched: the only failures are 12test_unary_spurious_fpexceptioncases forsign, which fail identically onmainand are unrelated to this change.