VQSelect/VQSort on aarch64/NEON: 2–4× slower than scalar on near-sorted input

[dmatth1]

Follow-up from the NaN work in #3098 / #3129:

I'm currently working on wiring VQSelect into NumPy's np.partition for aarch64/ppc64 (numpy/numpy#31506) and noticed that for near-sorted input VQSelect is slower than NumPy's scalar introselect. We are discussing adding a gate which i have prototyped (only call VQSelect when not nearsorted) for Numpy but we are wondering if a) this gate would be better suited to live in Highway and b) perhaps there is a better way to address the regression for Neon architectures instead of a gate/probe?

Measurement

M1 Pro (NEON), asv bench_function_base.Partition.time_partition, ARRAY_SIZE=100000. Ratio = VQSelect / scalar-introselect; <1 = VQSelect faster, >1 = slower.

dtype	random	reversed	uniform	ordered	sb/10	sb/100	sb/1000
int16	0.09	0.08	0.07	0.61	0.18	0.16	0.20–0.23
int32	0.34	0.12	0.14	1.17	0.61	—	0.47–0.86
int64	0.57	0.20	0.25	2.01–2.07	1.05–1.08	1.66–1.76	0.81–1.42
float32	0.45	0.25	0.35	2.71–2.74	0.79–0.84	1.73–1.78	0.78–1.85
float64	0.70	0.38	0.40	4.00–4.04	1.34–1.42	2.55–2.58	1.24–2.68

(sb/N = sorted_block, sorted runs of length N.) Big wins on disordered input, but float64 ordered is 4× slower and sorted_block up to 2.6× slower; float32/int64 show the same shape.

Why it looks architectural

On NEON, the partition hot loop (StoreLeftRight → CompressStore) has no hardware compress. Highway emulates it with a shuffle-index table + byte lookup (arm_neon-inl.h) and NEON lacks a movemask too. NEON is also 128-bit (Lanes(double) == 2 vs 8 on AVX-512). So the vector partition does more work per element than on a hardware-compress, wide-vector target, enough that scalar introselect wins on near-sorted input. (I haven't isolated width vs compress-emulation)

I don't see how to close NEON's raw throughput. It can't be widened, Apple Silicon has no SVE, and Highway's vqsort disables itself on scalable SVE anyway (only fixed 128-bit SVE2_128 stays vectorized, per shared-inl.h). But may be wrong...

[johnplatts]

@dmatth1 and @jan-wassenberg - I raised the issue of whether a new PartitioningCompress op should be added in issue #3142 as there are targets such as PPC10, SVE, and RVV where there are performance benefits to having a separate PartitioningCompress op that has partitioning guarantees (where a generic compress-table implementation of PartitioningCompress is more efficient on PPC10, where combining the results of Compress(v, mask) and CompressNot(v, mask) can be efficiently done using detail::Splice on SVE, and where combining the results of Compress(v, mask) and CompressNot(v, mask) can be efficiently done using detail::SlideUp on RVV).

[jan-wassenberg]

Yes, vqsort is unfortunately much less helpful on NEON due to narrow vectors and lack of compress/movemask as you mention.
Do you have data on how many of your target machines support SVE2? That at least adds compress.
It might appear that vqsort disables itself for SVE, but that is due to compiler bugs relating to scalable vectors which may already be resolved, and the most common SVE2_128 target is anyway unaffected.

Thanks @johnplatts, your proposal could further improve perf on SVE and RVV.

For this issue, I think a "nearly sorted" sampling step is a good path forward. We haven't had or noticed any use cases where that would be helpful. If we integrate into vqsort, I'd like to minimize the overhead. We could possibly sample ~4 vectors, count the number of inversions, and if very low, have a closer look - is that similar to what you prototyped?

[dmatth1]

Yeah that's basically what I prototyped, a cheap sample to route and an exact closer look when it reads as near-sorted (which also handles the already-partitioned no-op).

One NEON-specific caveat from my testing: the near-sorted shapes (ordered, sorted_block) are caught reliably even by 4 vectors but on NEON a vector is only 2–4 lanes so 4 vectors is 8–16 elements. Too few for random input and it misroutes to scalar ~20–40% of the time, forfeiting the Highway win. 16 vectors makes random reliable on NEON (with basically no overhead).

Should I go ahead with raising a PR for the gate so we can review what I have prototyped?

[jan-wassenberg]

Sure, happy to review.
One thing to consider first: DrawSamples already does something similar, it draws 6x64 bytes and concentrates them into 2x64 using MedianOf3. That's exactly your desired size, I think. You could integrate your check there, perhaps also gated on HWY_MIN_BYTES <= 16? (because wider vectors are very likely still preferable)