Fast complex absolute argmax in Cython

Question

I'm thinking I implemented it optimally, but somehow it's much slower than what should be much slower, np.argmax(np.abs(x)). Where am I off?

Code rationale & results

• Mathematically, abs is sqrt(real**2 + imag**2), but argmax(abs(x)) == argmax(abs(x)**2), so no need for square root
• np.abs(x) also allocates and writes an array. Instead I overwrite a single value, current_abs2, which should eliminate allocation and only leave writing
• Argmax logic should be identical to NumPy's (I've not checked but only one best way to do it?)
• Views (R, I) are for... I don't recall, saw somewhere

So savings are in dropping sqrt and len(x)-sized allocation. Yet it's much slower...

%timeit np.argmax(np.abs(x))
%timeit abs_argmax(x.real, x.imag)

409  µs ± 2.33 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
3.09 ms ± 14.9 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

Here's the generated C code, just the function; the whole _optimized.c is 26000 lines.

The following Numba achieves 108 µs, very satisfactory, though I'm interested in why Cython fails.

Code

import cython

@cython.boundscheck(False)
@cython.wraparound(False)
cpdef int abs_argmax(double[:] re, double[:] im):
    # initialize variables
    cdef Py_ssize_t N = re.shape[0]
    cdef double[:] R = re  # view
    cdef double[:] I = im  # view

    cdef Py_ssize_t i = 0
    cdef int max_idx = 0
    cdef double current_max = 0
    cdef double current_abs2 = 0

    # main loop
    while i < N:
        current_abs2 = R[i]**2 + I[i]**2
        if current_abs2 > current_max:
            max_idx = i
            current_max = current_abs2
        i += 1

    # return
    return max_idx

Setup & execution

I use python setup.py build_ext --inplace, setup.py shown at bottom. Then,

import numpy as np
from _optimized import abs_argmax

x = np.random.randn(100000) + 1j*np.random.randn(100000)
%timeit np.argmax(np.abs(x))
%timeit abs_argmax(x.real, x.imag)

setup.py (I forget the rationale, just took certain recommendations)

from distutils import _msvccompiler
_msvccompiler.PLAT_TO_VCVARS['win-amd64'] = 'amd64'

from setuptools import setup, Extension
from Cython.Build import cythonize
import numpy as np

setup(
    ext_modules=cythonize(Extension("_optimized", ["_optimized.pyx"]),
                          language_level=3),
    include_dirs=[np.get_include()],
)

Environment

Windows 11, i7-13700HX CPU, Python 3.11.4, Cython 3.0.0, setuptools 68.0.0, numpy 1.24.4

Answer 1

In a generated code (slightly edited for readability)

__pyx_v_current_abs2 = (
    pow((*((double *) ((__pyx_v_R.data + __pyx_t_2 * __pyx_v_R.strides[0]) ))), 2.0) +
    pow((*((double *) ((__pyx_v_I.data + __pyx_t_3 * __pyx_v_I.strides[0]) ))), 2.0)
);

I do not like calls to pow. Apparently, cython is not smart enough, and transpiles ** 2 into a function call, rather that a simple multiplication. Try to help it:ㅤㅤㅤㅤㅤㅤ

    current_abs2 = R[i]*R[i] + I[i]*I[i]

and see what happens. As of the rest - failed branch predictions, missed vectorization, etc - we may only theorize.

Answer 2

The OP code suffers from at least a 6x factor of suckage. The question would benefit from posting godbolt links which examine the differences in the generated code.

Often the trouble boils down to

• inability to vectorize
• poor branch prediction

        if current_abs2 > current_max:

It seems plausible that branch prediction is falling apart there.

Consider pre-computing a bunch of abs2 values, sorting them descending, and returning the zeroth element. It's O(N log N), so theoretically worse by a factor of log N. But it might interact with memory hierarchy better in practice. Bench, and report back.

Answer 3

Thanks to vnp's answer; here's np.argmax(np.abs(x)) vs Numba vs Cython:

414  µs ± 8.09 µs per loop (mean ± std. dev. of 7 runs, 1,000  loops each)
89.9 µs ± 1.2  µs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
66.8 µs ± 1.11 µs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

Code

import cython

@cython.boundscheck(False)
@cython.wraparound(False)
cpdef int abs_argmax(double[:] R, double[:] I):
    # initialize variables
    cdef Py_ssize_t N = R.shape[0]

    cdef Py_ssize_t i = 0
    cdef int max_idx = 0
    cdef double current_max = 0
    cdef double current_abs2 = 0

    # main loop
    for i in range(N):
        current_abs2 = R[i]*R[i] + I[i]*I[i]
        if current_abs2 > current_max:
            max_idx = i
            current_max = current_abs2

    return max_idx