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

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)

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

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