diff --git a/CHANGELOG.md b/CHANGELOG.md index f0da8acf..0f9a81c6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,6 +13,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 c1a41660..16aea16f 100644 --- a/mkl_umath/src/mkl_umath_loops.c.src +++ b/mkl_umath/src/mkl_umath_loops.c.src @@ -1306,6 +1306,47 @@ mkl_umath_@TYPE@_@kind@(char **args, const npy_intp *dimensions, const npy_intp *oi @OP@= ri; 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]; + } + } + } 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..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):