From acced6be7c80b52786574718e1ccdf2dbec49476 Mon Sep 17 00:00:00 2001 From: Python Infra CI Date: Mon, 5 Oct 2026 18:33:58 +0000 Subject: [PATCH 1/3] perf: add contiguous fast path to complex add/subtract fallback loop --- mkl_umath/src/mkl_umath_loops.c.src | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/mkl_umath/src/mkl_umath_loops.c.src b/mkl_umath/src/mkl_umath_loops.c.src index c1a41660..9ffb132a 100644 --- a/mkl_umath/src/mkl_umath_loops.c.src +++ b/mkl_umath/src/mkl_umath_loops.c.src @@ -1306,6 +1306,28 @@ mkl_umath_@TYPE@_@kind@(char **args, const npy_intp *dimensions, const npy_intp *oi @OP@= ri; return; } + else if (contig) { + /* + * Inputs/output are contiguous (though possibly too small for + * the VML call above, or overlapping). Real and imaginary + * parts can be processed as one flat run of 2*n @ftype@ + * elements since +/- is applied componentwise in the same + * element order as the generic strided loop below, which + * lets the compiler autovectorize instead of walking the + * array through char* pointer arithmetic one complex value + * at a time. + */ + const @ftype@ *ip1 = (@ftype@ *)args[0]; + const @ftype@ *ip2 = (@ftype@ *)args[1]; + @ftype@ *op1 = (@ftype@ *)args[2]; + const npy_intp n2 = dimensions[0] * 2; + npy_intp i; + + NPY_PRAGMA_VECTOR + for (i = 0; i < n2; i++) { + op1[i] = ip1[i] @OP@ ip2[i]; + } + } else { BINARY_LOOP { const @ftype@ in1r = ((@ftype@ *)ip1)[0]; From 007298aad3a2da70594a6af4a2975ed2d8754638 Mon Sep 17 00:00:00 2001 From: "Harlow, Jordan" Date: Fri, 9 Oct 2026 12:13:55 -0600 Subject: [PATCH 2/3] Address review: gate complex add/subtract fast path on can_vectorize - Take the flat loop only when operands are disjoint or exactly in-place; accumulate (out == in1 + 1) goes back to the strided loop. Under GCC, NPY_PRAGMA_VECTOR (ivdep) gave wrong accumulate results. - Drop the pragma; it did not change icx codegen and added two -Wpass-failed warnings for CDOUBLE. - Give out == in1 and out == in2 their own loops so the runtime alias check does not send in-place calls to the scalar fallback. - Add tests for the contiguous and accumulate paths, and a changelog entry. --- CHANGELOG.md | 1 + mkl_umath/src/mkl_umath_loops.c.src | 49 ++++++++++++++++++++--------- mkl_umath/tests/test_basic.py | 42 +++++++++++++++++++++++++ 3 files changed, 77 insertions(+), 15 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 14f0aa29..9a2a0f9e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Changed * Raised the minimum build-time `Cython` requirement to `3.1.0`, the first release providing the `freethreading_compatible` directive [gh-255](https://github.com/IntelPython/mkl_umath/pull/255) * Pinned Cython in the Coverity Scan workflow so generated code stays stable between scans, and added `coverity/README.md` documenting the known Cython-boilerplate false positives and the scan review checklist [gh-266](https://github.com/IntelPython/mkl_umath/pull/266) +* Improved performance of `add` and `subtract` for contiguous `complex64` and `complex128` arrays below the oneMKL VM threshold, including in-place calls [gh-278](https://github.com/IntelPython/mkl_umath/pull/278) ### Fixed * Fixed `absolute` (float32/float64) returning `-NaN` for `+NaN` input on the scalar fallback path (same issue as numpy/numpy@dc478c58b9, gh-31433) [gh-269](https://github.com/IntelPython/mkl_umath/pull/269) diff --git a/mkl_umath/src/mkl_umath_loops.c.src b/mkl_umath/src/mkl_umath_loops.c.src index 9ffb132a..16aea16f 100644 --- a/mkl_umath/src/mkl_umath_loops.c.src +++ b/mkl_umath/src/mkl_umath_loops.c.src @@ -1306,26 +1306,45 @@ mkl_umath_@TYPE@_@kind@(char **args, const npy_intp *dimensions, const npy_intp *oi @OP@= ri; return; } - else if (contig) { + else if (can_vectorize) { /* - * Inputs/output are contiguous (though possibly too small for - * the VML call above, or overlapping). Real and imaginary - * parts can be processed as one flat run of 2*n @ftype@ - * elements since +/- is applied componentwise in the same - * element order as the generic strided loop below, which - * lets the compiler autovectorize instead of walking the - * array through char* pointer arithmetic one complex value - * at a time. + * Contiguous and either disjoint or exactly in-place, but too + * small for the VML call above. +/- acts on the real and + * imaginary parts independently, so treat them as one flat + * run of 2*n @ftype@ elements. The in-place cases get their + * own loops so the compiler's runtime alias check compares + * only the two distinct buffers; with three pointers it sees + * out == in as an overlap and runs the scalar fallback. + * Partially overlapping operands (e.g. accumulate, where + * out == in1 + 1) take the strided loop below. */ - const @ftype@ *ip1 = (@ftype@ *)args[0]; - const @ftype@ *ip2 = (@ftype@ *)args[1]; - @ftype@ *op1 = (@ftype@ *)args[2]; const npy_intp n2 = dimensions[0] * 2; npy_intp i; - NPY_PRAGMA_VECTOR - for (i = 0; i < n2; i++) { - op1[i] = ip1[i] @OP@ ip2[i]; + if (args[2] == args[0]) { + @ftype@ *iop1 = (@ftype@ *)args[0]; + const @ftype@ *ip2 = (const @ftype@ *)args[1]; + + for (i = 0; i < n2; i++) { + iop1[i] = iop1[i] @OP@ ip2[i]; + } + } + else if (args[2] == args[1]) { + const @ftype@ *ip1 = (const @ftype@ *)args[0]; + @ftype@ *iop2 = (@ftype@ *)args[1]; + + for (i = 0; i < n2; i++) { + iop2[i] = ip1[i] @OP@ iop2[i]; + } + } + else { + const @ftype@ *ip1 = (const @ftype@ *)args[0]; + const @ftype@ *ip2 = (const @ftype@ *)args[1]; + @ftype@ *op1 = (@ftype@ *)args[2]; + + for (i = 0; i < n2; i++) { + op1[i] = ip1[i] @OP@ ip2[i]; + } } } else { diff --git a/mkl_umath/tests/test_basic.py b/mkl_umath/tests/test_basic.py index 8433428c..6db3dd5a 100644 --- a/mkl_umath/tests/test_basic.py +++ b/mkl_umath/tests/test_basic.py @@ -25,6 +25,7 @@ import numpy as np import pytest +from numpy.testing import assert_array_equal import mkl_umath._ufuncs as mu @@ -191,6 +192,47 @@ def test_reduce_complex(func, dtype): ), f"Results for '{func}[reduce]' do not match" +def _complex_array(size, dtype): + re = np.random.uniform(-10, 10, size) + im = np.random.uniform(-10, 10, size) + return (re + 1j * im).astype(dtype) + + +@pytest.mark.parametrize("func", ["add", "subtract"]) +@pytest.mark.parametrize("size", [1, 7, 100, 1001]) +@pytest.mark.parametrize("dtype", [np.complex64, np.complex128]) +def test_complex_contig(func, size, dtype): + # testing the contiguous branch below the VML threshold: out-of-place + # and in-place on either input; +/- is exact, so compare bitwise + a = _complex_array(size, dtype) + b = _complex_array(size, dtype) + mkl_func = getattr(mu, func) + np_res = getattr(np, func)(a, b) + + out = np.empty_like(a) + mkl_func(a, b, out=out) + assert_array_equal(out, np_res) + + a_inplace = a.copy() + mkl_func(a_inplace, b, out=a_inplace) + assert_array_equal(a_inplace, np_res) + + b_inplace = b.copy() + mkl_func(a, b_inplace, out=b_inplace) + assert_array_equal(b_inplace, np_res) + + +@pytest.mark.parametrize("func", ["add", "subtract"]) +@pytest.mark.parametrize("dtype", [np.complex64, np.complex128]) +def test_accumulate_complex(func, dtype): + # accumulate calls the loop with out == in1 + 1 element, so each output + # depends on the previous one and must take the strided branch + a = _complex_array(100, dtype) + mkl_res = getattr(mu, func).accumulate(a) + np_res = getattr(np, func).accumulate(a) + assert_array_equal(mkl_res, np_res) + + @pytest.mark.parametrize("size", [100, 8192 + 1]) @pytest.mark.parametrize("dtype", [np.float32, np.float64]) def test_absolute_nan_signbit(size, dtype): From e5bf8d2d1983033d5e2c7f1600453c67a6e1054b Mon Sep 17 00:00:00 2001 From: "Harlow, Jordan" Date: Mon, 19 Oct 2026 11:33:48 -0600 Subject: [PATCH 3/3] Address review: move the contiguous complex add/subtract loops out of line, add a z += z loop and tests --- mkl_umath/src/mkl_umath_loops.c.src | 98 +++++++++++++++++------------ mkl_umath/tests/test_basic.py | 45 ++++++++++++- 2 files changed, 102 insertions(+), 41 deletions(-) diff --git a/mkl_umath/src/mkl_umath_loops.c.src b/mkl_umath/src/mkl_umath_loops.c.src index 16aea16f..ad517e76 100644 --- a/mkl_umath/src/mkl_umath_loops.c.src +++ b/mkl_umath/src/mkl_umath_loops.c.src @@ -1282,6 +1282,64 @@ pairwise_sum_@TYPE@(@ftype@ *rr, @ftype@ * ri, char * a, npy_intp n, npy_intp st * #PW = 1, 0# * #VML = Add, Sub# */ +/* + * Kept out of line so that adding these loops does not change the + * register allocation of the strided BINARY_LOOP in the caller. + */ +NPY_NOINLINE void +contig_@kind@_@TYPE@(char **args, npy_intp n) +{ + /* + * Contiguous and either disjoint or exactly in-place, but too + * small for the VML call in the caller. +/- acts on the real and + * imaginary parts independently, so treat them as one flat + * run of 2*n @ftype@ elements. The in-place cases get their + * own loops so icx's runtime alias check compares only the two + * distinct buffers; with three pointers icx sees out == in as + * an overlap and runs its scalar fallback (gcc and clang + * vectorize the three-pointer loop either way). + * Partially overlapping operands (e.g. accumulate, where + * out == in1 + 1) take the caller's strided loop. + */ + const npy_intp n2 = n * 2; + npy_intp i; + + if (args[2] == args[0]) { + @ftype@ *iop1 = (@ftype@ *)args[0]; + + if (args[1] == args[0]) { + /* z += z: a single pointer leaves no alias check to fail */ + for (i = 0; i < n2; i++) { + iop1[i] = iop1[i] @OP@ iop1[i]; + } + } + else { + const @ftype@ *ip2 = (const @ftype@ *)args[1]; + + for (i = 0; i < n2; i++) { + iop1[i] = iop1[i] @OP@ ip2[i]; + } + } + } + else if (args[2] == args[1]) { + const @ftype@ *ip1 = (const @ftype@ *)args[0]; + @ftype@ *iop2 = (@ftype@ *)args[1]; + + for (i = 0; i < n2; i++) { + iop2[i] = ip1[i] @OP@ iop2[i]; + } + } + else { + const @ftype@ *ip1 = (const @ftype@ *)args[0]; + const @ftype@ *ip2 = (const @ftype@ *)args[1]; + @ftype@ *op1 = (@ftype@ *)args[2]; + + for (i = 0; i < n2; i++) { + op1[i] = ip1[i] @OP@ ip2[i]; + } + } +} + void mkl_umath_@TYPE@_@kind@(char **args, const npy_intp *dimensions, const npy_intp *steps, void *NPY_UNUSED(func)) { @@ -1307,45 +1365,7 @@ mkl_umath_@TYPE@_@kind@(char **args, const npy_intp *dimensions, const npy_intp return; } else if (can_vectorize) { - /* - * Contiguous and either disjoint or exactly in-place, but too - * small for the VML call above. +/- acts on the real and - * imaginary parts independently, so treat them as one flat - * run of 2*n @ftype@ elements. The in-place cases get their - * own loops so the compiler's runtime alias check compares - * only the two distinct buffers; with three pointers it sees - * out == in as an overlap and runs the scalar fallback. - * Partially overlapping operands (e.g. accumulate, where - * out == in1 + 1) take the strided loop below. - */ - const npy_intp n2 = dimensions[0] * 2; - npy_intp i; - - if (args[2] == args[0]) { - @ftype@ *iop1 = (@ftype@ *)args[0]; - const @ftype@ *ip2 = (const @ftype@ *)args[1]; - - for (i = 0; i < n2; i++) { - iop1[i] = iop1[i] @OP@ ip2[i]; - } - } - else if (args[2] == args[1]) { - const @ftype@ *ip1 = (const @ftype@ *)args[0]; - @ftype@ *iop2 = (@ftype@ *)args[1]; - - for (i = 0; i < n2; i++) { - iop2[i] = ip1[i] @OP@ iop2[i]; - } - } - else { - const @ftype@ *ip1 = (const @ftype@ *)args[0]; - const @ftype@ *ip2 = (const @ftype@ *)args[1]; - @ftype@ *op1 = (@ftype@ *)args[2]; - - for (i = 0; i < n2; i++) { - op1[i] = ip1[i] @OP@ ip2[i]; - } - } + contig_@kind@_@TYPE@(args, dimensions[0]); } else { BINARY_LOOP { diff --git a/mkl_umath/tests/test_basic.py b/mkl_umath/tests/test_basic.py index 6db3dd5a..0805a216 100644 --- a/mkl_umath/tests/test_basic.py +++ b/mkl_umath/tests/test_basic.py @@ -199,11 +199,11 @@ def _complex_array(size, dtype): @pytest.mark.parametrize("func", ["add", "subtract"]) -@pytest.mark.parametrize("size", [1, 7, 100, 1001]) +@pytest.mark.parametrize("size", [1, 7, 100, 1001, 10**5]) @pytest.mark.parametrize("dtype", [np.complex64, np.complex128]) def test_complex_contig(func, size, dtype): # testing the contiguous branch below the VML threshold: out-of-place - # and in-place on either input; +/- is exact, so compare bitwise + # and in-place on either input; results must equal NumPy's exactly a = _complex_array(size, dtype) b = _complex_array(size, dtype) mkl_func = getattr(mu, func) @@ -233,6 +233,47 @@ def test_accumulate_complex(func, dtype): assert_array_equal(mkl_res, np_res) +@pytest.mark.parametrize("func", ["add", "subtract"]) +@pytest.mark.parametrize("size", [1, 7, 100, 1001, 10**5]) +@pytest.mark.parametrize("dtype", [np.complex64, np.complex128]) +def test_complex_contig_same_input(func, size, dtype): + # in1 is in2: out-of-place takes the disjoint loop with ip1 == ip2, + # out=a takes the out == in1 loop with ip2 == iop1 + a = _complex_array(size, dtype) + mkl_func = getattr(mu, func) + np_res = getattr(np, func)(a, a) + + assert_array_equal(mkl_func(a, a), np_res) + + a_inplace = a.copy() + mkl_func(a_inplace, a_inplace, out=a_inplace) + assert_array_equal(a_inplace, np_res) + + +@pytest.mark.parametrize("func", ["add", "subtract"]) +@pytest.mark.parametrize("dtype", [np.complex64, np.complex128]) +def test_complex_contig_signed_zero(func, dtype): + # every (+-0, +-0) pair in both the real and the imaginary part; + # assert_array_equal treats -0.0 == 0.0, so compare the sign bits + ftype = np.float32 if dtype == np.complex64 else np.float64 + zeros = np.array([0.0, -0.0], dtype=ftype) + x, y = np.meshgrid(zeros, zeros) + a = np.tile(np.repeat(x.ravel(), 2), 8).view(dtype) + b = np.tile(np.repeat(y.ravel(), 2), 8).view(dtype) + mkl_func = getattr(mu, func) + np_res = getattr(np, func)(a, b) + + for out_is in ("new", "a", "b"): + a_, b_ = a.copy(), b.copy() + out = {"new": np.empty_like(a), "a": a_, "b": b_}[out_is] + mkl_func(a_, b_, out=out) + assert_array_equal( + np.signbit(out.view(ftype)), + np.signbit(np_res.view(ftype)), + err_msg=f"signed zero mismatch for '{func}' with out={out_is}", + ) + + @pytest.mark.parametrize("size", [100, 8192 + 1]) @pytest.mark.parametrize("dtype", [np.float32, np.float64]) def test_absolute_nan_signbit(size, dtype):