perf: add contiguous fast path to complex add/subtract fallback loop

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)

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];

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):