diff --git a/CHANGELOG.md b/CHANGELOG.md index 2e23b390..98ad4cc3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -21,6 +21,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 3d0f2fcb..d1b335ca 100644 --- a/mkl_umath/src/mkl_umath_loops.c.src +++ b/mkl_umath/src/mkl_umath_loops.c.src @@ -1244,6 +1244,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)) { @@ -1268,6 +1326,9 @@ mkl_umath_@TYPE@_@kind@(char **args, const npy_intp *dimensions, const npy_intp *oi @OP@= ri; return; } + else if (can_vectorize) { + contig_@kind@_@TYPE@(args, dimensions[0]); + } else { BINARY_LOOP { const @ftype@ in1r = ((@ftype@ *)ip1)[0]; diff --git a/mkl_umath/tests/test_basic.py b/mkl_umath/tests/test_basic.py index 8433428c..0805a216 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,88 @@ 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, 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; results must equal NumPy's exactly + 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("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):