ENH: use x86-simd-sort for descending sorts and partitions - #32690

Measured on AVX2 over 1M elements, descending now runs at the speed of ascending: ~4-9x faster than before for sort/partition and ~2-3x for argsort/argpartition.

Is there any chance you could share the benchmark? I'd be interested to run it on my ryzen cpu that supports avx512 as well.

Measured on AVX2 over 1M elements, descending now runs at the speed of ascending: ~4-9x faster than before for sort/partition and ~2-3x for argsort/argpartition.

Is there any chance you could share the benchmark? I'd be interested to run it on my ryzen cpu that supports avx512 as well.

Results on my machine: i7-13650HX (AVX2, no AVX-512), Linux, CPython 3.14, release builds. The number from the OP is on a different machine.

"""
pyperf benchmark for numpy/numpy#32690: x86-simd-sort for descending
sort / argsort / partition / argpartition.

Run once per NumPy build and compare the results:

    python bench_descending_sort.py -o main.json      # in the env with main
    python bench_descending_sort.py -o pr.json        # in the env with the PR
    python -m pyperf compare_to main.json pr.json --table

Options (on top of the usual pyperf ones such as --fast, --rigorous,
--affinity):

    --dtypes float64,int32     dtypes to run (default: see DTYPES)
    --sizes 10000,1000000      array sizes to run (default: 1000000)
    --funcs sort,argsort       functions to run (default: all four)
    --kth-frac 0.5             kth for (arg)partition, as a fraction of n

All functions are called out-of-place (``np.sort(a)``, not ``a.sort()``) so
every call sees the same unsorted input.  That includes a copy of the input
(or the creation of an index array), which is the same on both builds; the
``copy`` benchmark measures that fixed cost so it can be subtracted.
"""
from functools import partial

import numpy as np
import pyperf

DTYPES = ["float64", "float32", "float16", "int64", "int32", "int16"]
SIZES = [1_000_000]
FUNCS = ["sort", "argsort", "partition", "argpartition"]


def make_data(dtype, n):
    rng = np.random.default_rng(12345)
    dtype = np.dtype(dtype)
    if dtype.kind == "f":
        return rng.standard_normal(n).astype(dtype)
    info = np.iinfo(dtype)
    return rng.integers(info.min, info.max, size=n, dtype=dtype, endpoint=True)


def add_cmdline_args(cmd, args):
    cmd.extend(["--dtypes", args.dtypes, "--sizes", args.sizes,
                "--funcs", args.funcs, "--kth-frac", str(args.kth_frac)])


def main():
    runner = pyperf.Runner(add_cmdline_args=add_cmdline_args)
    parser = runner.argparser
    parser.add_argument("--dtypes", default=",".join(DTYPES))
    parser.add_argument("--sizes", default=",".join(map(str, SIZES)))
    parser.add_argument("--funcs", default=",".join(FUNCS))
    parser.add_argument("--kth-frac", type=float, default=0.5)
    args = runner.parse_args()

    # Recorded in the JSON so it is clear which build and which SIMD
    # kernels produced the numbers (`pyperf metadata file.json`).
    from numpy._core._multiarray_umath import __cpu_features__
    runner.metadata["numpy_version"] = np.__version__
    runner.metadata["numpy_file"] = np.__file__
    runner.metadata["numpy_cpu_features"] = " ".join(
        name for name, have in __cpu_features__.items() if have)

    for dtype in args.dtypes.split(","):
        for n in map(int, args.sizes.split(",")):
            a = make_data(dtype, n)
            kth = min(n - 1, int(n * args.kth_frac))
            tag = f"{dtype} n={n}"

            runner.bench_func(f"copy {tag}", a.copy)
            for name in args.funcs.split(","):
                func = getattr(np, name)
                extra = (kth,) if "partition" in name else ()
                for order, descending in (("asc", False), ("desc", True)):
                    # bench_func() does not forward keyword arguments
                    runner.bench_func(
                        f"{name} {tag} {order}",
                        partial(func, a, *extra, descending=descending))


if __name__ == "__main__":
    main()

I'm not getting that significant speed up on AMD Ryzen 7 7800X3D. I ran the benchmark you posted earlier by installing numpy like python -m pip install -vvv . -C'setup-args=-Dbuildtype=release' -C'build-dir=build-dir' --no-clean --no-build-isolation on main and on your PR branch. Then I ran python -m pyperf compare_to --table --table-format md main.json branch.json to compare the pyperf jsons. And this is the output:

Benchmark hidden because not significant (21): sort float64 n=1000000 asc, partition float64 n=1000000 asc, argpartition float64 n=1000000 asc, copy float32 n=1000000, sort float32 n=1000000 asc, argsort float32 n=1000000 asc, partition float32 n=1000000 asc, copy float16 n=1000000, partition float16 n=1000000 asc, argpartition float16 n=1000000 desc, sort int64 n=1000000 asc, argsort int64 n=1000000 asc, partition int64 n=1000000 asc, copy int32 n=1000000, sort int32 n=1000000 asc, argsort int32 n=1000000 asc, partition int32 n=1000000 asc, argpartition int32 n=1000000 asc, sort int16 n=1000000 asc, partition int16 n=1000000 asc, argpartition int16 n=1000000 desc

I was about to check on another AVX 512 arch on a local supercomputer (Intel(R) Xeon(R) Platinum 8480+). I'll move on though, since Ryzen should do it.

I was about to check on another AVX 512 arch on a local supercomputer (Intel(R) Xeon(R) Platinum 8480+). I'll move on though, since Ryzen should do it.

yours supports AVX512_SPR as well though which mine doesn't so if you wanna do it, go ahead.

Okay yeah I ran the benchmarks with NPY_DISABLE_CPU_FEATURES="X86_V4 AVX512_ICL AVX512_SPR" and apparently it's just a known thing that Zen 4 sorting with avx512 is just bad. I'm getting this below on avx2 with my cpu as well. I'm wondering if we can do something to automatically disable this for people on Zen 4 🤣. Dispatching to avx512 on Zen 4 feels wrong.

My AI model informs me that:

• On Intel (Skylake-X, Ice Lake, Sapphire Rapids), x86-simd-sort's AVX-512 path is its fastest path, faster than AVX2. That's the hardware it was written and tuned for, at Intel.
• On Zen 4, most AVX-512 instructions run fine, but vpcompressd/q/ps/pd with a memory destination is microcoded and very slow, reportedly over 100 cycles per instruction. The register-destination form of the same instruction is fast. x86-simd-sort's partition step does exactly this memory-destination compress-store on every vector (_mm512_mask_compressstoreu_*), so it hits the worst case.
• Zen 5 reportedly fixed this.

Edit: apparently this is know and there are issues already open for it

So LGTM too! Not formally approving as I very quickly skimmed through the code but it looked like simple changes.