ENH: Use Highway's VQSort on AArch64#24018
Conversation
| for sort_file in [ | ||
| "vqsort.cc", | ||
| "vqsort_128a.cc", | ||
| "vqsort_128d.cc", | ||
| "vqsort_f32a.cc", | ||
| "vqsort_f32d.cc", | ||
| "vqsort_f64a.cc", | ||
| "vqsort_f64d.cc", | ||
| "vqsort_i16a.cc", | ||
| "vqsort_i16d.cc", | ||
| "vqsort_i32a.cc", | ||
| "vqsort_i32d.cc", | ||
| "vqsort_i64a.cc", | ||
| "vqsort_i64d.cc", | ||
| "vqsort_kv64a.cc", | ||
| "vqsort_kv64d.cc", | ||
| "vqsort_kv128a.cc", | ||
| "vqsort_kv128d.cc", | ||
| "vqsort_u16a.cc", | ||
| "vqsort_u16d.cc", | ||
| "vqsort_u32a.cc", | ||
| "vqsort_u32d.cc", | ||
| "vqsort_u64a.cc", | ||
| "vqsort_u64d.cc", | ||
| ] |
There was a problem hiding this comment.
This isn't ideal, when compiling without these files I got missing symbol issues, but we definitely don't use them all and they add ~100M to the .so.
@jan-wassenberg any thoughts on how to minimize the number of files we need here?
There was a problem hiding this comment.
I think we're only using {u/i/f}{32/64}a in your dispatcher and we ought to be able to remove 128*, 16*, and *d.cc (the latter for descending order).
If there are still linker errors after removing only those, then it's likely due to the dllexport. What we could then do is add a macro bitfield that you could define to replace the implementation of the unused/unwanted types with a stub.
Something like VQSORT_TYPES which contains multiple groups (one each for uint/int/float/KV), where each group has one bit encoding whether the element size equal to the bit's value is enabled.
Example, for only say unsigned/float types: #define VQSORT_TYPES (((4 | 8) << 6) | (4 | 8)).
Would that be useful?
There was a problem hiding this comment.
As an example, when I try to run tests:
Original error was: /numpy/build/testenv/lib/python3.10/site-packages/numpy/core/_multiarray_umath.cpython-310-aarch64-linux-gnu.so: undefined symbol: _ZN3hwy6VQSortEPlmNS_14SortDescendingE
The other option is I could use vqsort-inl.h and copy the implementations across from the various vqsort_X.cc files? That'd also mean I can do float16 without adding a vsqort_f16a.cc.
There was a problem hiding this comment.
Could you add something like this to vqsort-inl.h ? I notice there's a lot of detail which would be bad to include in the NumPy source.
template <typename T>
void VQSort(T* HWY_RESTRICT keys, size_t n, SortAscending) {
SortTag<T> d;
detail::SharedTraits<detail::TraitsLane<detail::OrderAscending<T>>> st;
Sort(d, st, keys, num);
}
If not the bitfield approach is also good 😸
There was a problem hiding this comment.
Sure, we can add that. As written, this will not support the 128-bit nor key-value types - is that OK?
There was a problem hiding this comment.
Sure, we can add that. As written, this will not support the 128-bit nor key-value types - is that OK?
That'd be specialising around uint128_t with the same uint64*2 logic as in the .cc file? I think we can add it later if necessary, there's no sort dispatch for it atm.
There was a problem hiding this comment.
Right. And KV types also use different Traits.
Sounds good, we can indeed add that later.
| @@ -0,0 +1,44 @@ | |||
| /*@targets | |||
There was a problem hiding this comment.
I split this out as Highway's sort doesn't have an implementation of argsort, unless I'm wrong @jan-wassenberg?
There was a problem hiding this comment.
FYI argsort can be implemented as copying an array to hwy::K32V32 or K64V64 types (putting the key in .key and 0..N-1 into .value) and then copying only the .value back to the array.
PiperOrigin-RevId: 542571742
PiperOrigin-RevId: 542571742
PiperOrigin-RevId: 542571742
PiperOrigin-RevId: 542579202
4a6edf9 to
50bd401
Compare
I think this is worthy of wider discussion. So far we have
Perhaps we should stop and think where this all is leading. Maybe NumPy should not be in the business of providing highly specialized inner loop functions, and we should instead provide an easily pluggable interface for others to experiment with, together with a base implementation and test cases for accuracy (known good answers) / precision (same answers across various array shapes). |
Yip, I've been struggling with the current set of NumPy variations on the theme of performance optimisation, it's not clear what the current best practice is or where it's headed. Landing these vendored libraries (such as My interpretation of the proposal for using SLEEF (alongside my reasoning for backing something like this) is that people want NumPy to perform better out-of-the-box and SLEEF already has a bunch of cross-platform routines for this; indicating that there is a desire for things to be faster, but these efforts haven't succeeded so far. That demand helps me to understand that we should be taking a more off-the-shelf solution, which means doing less work in NumPy but still meeting the original goals of the universal intrinsics work - which is what I see when I stumble across Highway. We could take a lot of complexity out of NumPy itself by leveraging the Highway's dynamic dispatch, optimised routines and existing universal intrinsics solution (which already supports scalable vectors). I don't think performance should be passed back to users as plug-ins, if I wanted performance without using NumPy itself, I could easily move to something like Numba, one of the other optimisation libraries or stop using NumPy (as the de facto standard, such a stance seems rather out of place). I did like the errstate proposal, which allows users to opt into the lower accuracy, higher throughput routines but the high/low accuracy problem is only applicable to specific routines. Returning to this case, it's a sort routine, if it puts things in the wrong order it's wrong - these optimisations can benefit all users without much controversy which seems like something everyone would want? Please don't lose sight of the benefits this specific PR brings whilst considering whether to integrate more of Highway in the future. |
|
We discussed this a bit yesterday, please correct me if I am wrong (or just disagree), but basically. The way I interpreted this, there was also a hope to use highway more generally. Sentiment was a bit towards:
Overall, there wasn't much enthusiasm for this yesterday. |
|
Hi, Highway/vqsort main author here :)
Not sure this is the case. Our intrinsics are just fully-inlined wrapper functions. Only the sort implementation itself would contribute to binary size. VQSort also has some advantages concerning speed on AVX-512 (see https://github.com/Voultapher/sort-research-rs/blob/main/writeup/intel_avx512/text.md#updated-vqsort-results) and also robustness vs skewed input distributions, see other parts of that article. Happy to discuss further, including via video call, if you are interested. |
@jan-wassenberg , I think this would just be something like |
|
I think the larger point is that we are uncomfortable adding yet another submodule with inner loop implementations to NumPy. We would prefer to generalize the routines we already have for x86_64, using NumPy's universal intrinsics, to the other architectures. |
|
Maybe we can split this up a bit, since I probably munched it too much yesterday. I am not sure how much we gain by adopting highway more broadly, there were arguments against it, and maybe in part that is because Sayed has momentum for the "custom" universal intrinsics but there were more arguments. (I honestly would prefer to just steal something, but SIMD is already hairy with compiler bugs, and I don't know it well enough to judge whether we would just run into more issues.) So this is about sorting: I think right now it just switches things on AArch64? But highway supports also AVX512, so in the end it would really be swapping out one implementation for another. This doesn't actually stop us from reversing it out again once we can, so the question is if there is a downside? (if there are issues, too bad, we will revert) I probably misunderstood that this had some impact on binary size, I have to admit, if it doesn't I am not sure I see a reason for not swapping in highway for sorting (completely) until Raghuveer's work supports more broadly SIMD. |
@mattip, that seems to disadvantage those who didn't author the x86_64 code, whilst there is an equivalent solution available for many architectures. I'm personally reluctant to invest further effort in porting to universal intrinsics, there's solutions already available to users as proven by the SVML and x86-simd-sort integrations. Previous universal intrinsic work was already a wasted effort now that the math routines have been reverted and accuracy changed in SVML.
Where is this momentum? It'd be great to see things like #22265 landed but it has stalled waiting for the C++ intrinsic work. As far as compiler bugs, I think you'll get them either way, they're bugs after all 😸
Sorry @seberg, that's likely down to me, originally this was pretty bloated thanks to including many separate compilation units but through @jan-wassenberg's efforts on
Thanks for bringing this back to the benefits of this contribution itself, I did limit this to AArch64 to minimise the overall impact, mirroring the accepted x86-simd-sort integration - it wouldn't be hard to enable this for AVX512 as well and potentially further. One limitation is it doesn't quite drop in for FP16 (@jan-wassenberg can maybe advise), nor argsort but we can probably implement that ourselves based on the suggestion above. |
I am sorry you feel that way. I seem to remember having similar feelings when SVML was proposed, and that the goal all along was to replace SVML with universal intrinsics. Unfortunately we did not adequately address accuracy when adding the tests that were a condition for the SVML acceptance, and that led to the mess around sin/cos. At the end of the day mine is just one opinion of many maintainers. I think the place to push for a revamped approach to SIMD interinsics is discuss replacing NEP 38 on the mailing list, or to create a PR amending that NEP. |
|
There are two aspects to this PR:
(1) is tied with (2): you can't port sorting routines without using highway as the SIMD framework which duplicates NumPy universal intrinsics. If NumPy wants to use highway as the SIMD framework, that's a whole new discussion and as @mattip says it would require modifying NEP 38. If the concern is just about sorting routines on AArch64, we could address that by porting over x86-simd-sort using NumPy universal intrinsics once we merge #21057. If this is the direction we want to go, I will be happy to collaborate with @Mousius to port it. Note that x86-simd-sort also supports O(1) space |
I don't think we should be considering the implementation details of the sort algorithm, just as From a technical point of view, this is exactly the same as implementing sort using I'm also getting a distinct impression from these conversations that not many are in favour of replacing NEP38, as only @seberg has spoken in favour of using something already written.
I can't guarantee I'll be able to resource porting anything further, if/when C++ intrinsics are mature enough to do so, we'll have to re-assess then. It seems to be creating more work given VQSort already supports multiple platforms, but it is already the chosen solution so maybe we'll have to. |
|
I agree there are two discussions here; let's move the intrinsics part to a separate issue.
Yes, I happen to be working on fp16 at the moment :) It's a larger undertaking, so will take several days.
Understandable. With VQSort author hat on, I'd point out that it is also considerably faster and It might be possible to copy over some of the techniques from VQSort while porting, though. Have we estimated the number of new universal intrinsics that would be required? I imagine that more than a few new ones will be required for the shuffles used in sorting networks.
To adopting VQSort as suggested in this PR? Adding any code increases compile time and binary size, but I do not believe that is problematic here, so no, I do not see downsides for numpy. |
|
|
||
| # Set NPY_DISABLE_HIGHWAY=1 in the environment to disable the Highway | ||
| # library. This option only has significance on a AArch64 and is most | ||
| # useful to avoid improperly requiring Highway when cross compiling. |
There was a problem hiding this comment.
This seems copied from the SVML snippet, however the reason that is there for SVML is that that library gets shipped as object files rather than as source files. There shouldn't be an issue with Highway when cross-compiling I hope?
7243554 to
278e1de
Compare
Not sure about the complexity you're facing. My suggestion was similar to your current approach, but instead of creating a new source
I didn't mean to come off as negative at all. I just thought that adding a few extra CPU targets to your patch might not change much, maybe just a line or two. My conclusion was based on your response and my understanding that enabling a target like
Offering multiple choices should be considered if there's a performance benefit. If both solutions are sufficiently isolated, it shouldn't lead to a maintenance burden.
This is the reason I suggested leveraging the broad support and enabling additional CPU features of the architectures that are already supported by Highway.
Once again, I don't mean to be negative at all, and I will try to make my points more clear to avoid any kind of misunderstanding.
During our last optimization meeting, we agreed to accept your work as it is, without any further changes, to expedite progress. Later, we may focus on separating the two libraries. |
We can just agree to disagree here 😸, I've been out sick for several weeks and looking at a phased return, so actually getting to do the work is difficult. @r-devulap has expressed willingness to do the follow up which I really appreciated!
Aha, that would be my preference, and then I could focus on enabling AVX2 if we needed in a single patch - dependent on whether we agree that using Highway rather than x86-simd-sort is worth it.
Every branch in code or extra dependency is a maintenance burden, but I'm not going to worry too much about this one, given @r-devulap is keen to continue maintenance.
Yip, this comment was before the meeting. I believe we did reach a compromise where this patch would land, and @r-devulap would support me in doing a follow-up (which, again, I really appreciate!). I'm excited to see it land after such long discussions 😸 |
|
I would prefer to have #25045 for AVX2 versions of sort and partition. It will soon be followed up with AVX2 versions of argsort and argpartition. |
|
@Mousius, you are welcome! Once this is merged, I will update #24924 to split x86-simd-sort into its own dispatch module like how @seiko2plus prefers it. |
|
@Mousius, I think we should have a release note for this pr. |
np.partition (and callers: pandas nlargest, sklearn k-NN, etc.) on aarch64 has always fallen through to a 1998 scalar quickselect. Highway VQSelect (added to vqsort-inl.h in April 2024) provides a vectorized nth_element via NEON's vtbl-based Compress primitive. This change wires it into NumPy's partition dispatch alongside the existing x86-simd-sort path. Changes ------- numpy/_core/src/npysort/highway_qselect.dispatch.cpp (new) Dispatch file wrapping VQSelectStatic(arr, num, kth, SortAscending()). Covers int32, uint32, int64, uint64, float, double. Mirrors the shape of highway_qsort.dispatch.cpp. numpy/_core/src/npysort/highway_qsort.hpp Add QSelect dispatch declaration alongside QSort. numpy/_core/src/npysort/selection.cpp Add #if NPY_CPU_AMD64/X86 / #else split in quickselect_dispatch, matching the existing pattern in quicksort.hpp. x86 still uses x86-simd-sort; non-x86 now tries Highway VQSelect before falling back to scalar introselect. Adds HWY_QSELECT_MIN_N = 16384: below this threshold VQSelect's dispatch overhead dominates on L1-resident arrays (measured crossover ~16K elems). Removes dead DISABLE_HIGHWAY_OPTIMIZATION macro (defined but never used). numpy/_core/meson.build Register highway_qselect.dispatch.cpp for ASIMD and VSX2 targets. Scope ----- np.partition is fixed; np.argpartition is not (VQSelect operates on a keys-only array and cannot maintain a parallel index array). Benchmark (M1 Pro, Apple Silicon NEON, 40-rep median) ------------------------------------------------------ scalar = numpy 2.4.4 venv; vqselect = this branch dtype n k scalar vqselect speedup ------------------------------------------------------- float32 100K 10 5.0 2.3 ns/e 2.2x float32 1M 100 6.8 2.6 ns/e 2.6x float32 10M 100 8.9 2.7 ns/e 3.3x float32 10M 10K 6.3 2.7 ns/e 2.3x float64 100K 10 8.0 3.6 ns/e 2.2x float64 1M 100 9.2 4.5 ns/e 2.0x float64 10M 100 9.3 4.5 ns/e 2.1x float64 10M 10K 7.6 4.4 ns/e 1.7x int32 100K 10 6.3 1.3 ns/e 4.9x int32 1M 100 5.8 1.5 ns/e 3.8x int32 10M 100 4.2 1.7 ns/e 2.5x int64 1M 100 5.9 2.9 ns/e 2.0x int64 10M 100 5.8 3.1 ns/e 1.9x ------------------------------------------------------- n <= 16K (any dtype): threshold active, scalar path used, parity No regressions. np.sort on aarch64 already used Highway VQSort (PR numpy#24018); this patch brings the same acceleration to np.partition.
np.partition (and callers: pandas nlargest, sklearn k-NN, etc.) on aarch64 has always fallen through to a 1998 scalar quickselect. Highway VQSelect (added to vqsort-inl.h in April 2024) provides a vectorized nth_element via NEON's vtbl-based Compress primitive. This change wires it into NumPy's partition dispatch alongside the existing x86-simd-sort path. Changes ------- numpy/_core/src/npysort/highway_qselect.dispatch.cpp (new) Dispatch file wrapping VQSelectStatic(arr, num, kth, SortAscending()). Covers int32, uint32, int64, uint64, float, double. Mirrors the shape of highway_qsort.dispatch.cpp. numpy/_core/src/npysort/highway_qsort.hpp Add QSelect dispatch declaration alongside QSort. numpy/_core/src/npysort/selection.cpp Add #if NPY_CPU_AMD64/X86 / #else split in quickselect_dispatch, matching the existing pattern in quicksort.hpp. x86 still uses x86-simd-sort; non-x86 now tries Highway VQSelect before falling back to scalar introselect. Adds HWY_QSELECT_MIN_N = 16384: below this threshold VQSelect's dispatch overhead dominates on L1-resident arrays (measured crossover ~16K elems). Removes dead DISABLE_HIGHWAY_OPTIMIZATION macro (defined but never used). numpy/_core/meson.build Register highway_qselect.dispatch.cpp for ASIMD and VSX2 targets. Scope ----- np.partition is fixed; np.argpartition is not (VQSelect operates on a keys-only array and cannot maintain a parallel index array). Benchmark (M1 Pro, Apple Silicon NEON, 40-rep median) ------------------------------------------------------ scalar = numpy 2.4.4 venv; vqselect = this branch dtype n k scalar vqselect speedup ------------------------------------------------------- float32 100K 10 5.0 2.3 ns/e 2.2x float32 1M 100 6.8 2.6 ns/e 2.6x float32 10M 100 8.9 2.7 ns/e 3.3x float32 10M 10K 6.3 2.7 ns/e 2.3x float64 100K 10 8.0 3.6 ns/e 2.2x float64 1M 100 9.2 4.5 ns/e 2.0x float64 10M 100 9.3 4.5 ns/e 2.1x float64 10M 10K 7.6 4.4 ns/e 1.7x int32 100K 10 6.3 1.3 ns/e 4.9x int32 1M 100 5.8 1.5 ns/e 3.8x int32 10M 100 4.2 1.7 ns/e 2.5x int64 1M 100 5.9 2.9 ns/e 2.0x int64 10M 100 5.8 3.1 ns/e 1.9x ------------------------------------------------------- n <= 16K (any dtype): threshold active, scalar path used, parity No regressions. np.sort on aarch64 already used Highway VQSort (PR numpy#24018); this patch brings the same acceleration to np.partition.
np.partition on aarch64 has always fallen through to a 1998 scalar quickselect. np.sort already uses Highway VQSort (PR numpy#24018); this patch brings the same acceleration to np.partition via VQSelectStatic. Four files changed: - highway_qselect.dispatch.cpp: new dispatch file wrapping VQSelectStatic, covering int32/uint32/int64/uint64/float/double. Mirrors highway_qsort.dispatch.cpp. - highway_qsort.hpp: add QSelect dispatch declaration alongside QSort. - selection.cpp: add x86/non-x86 split in quickselect_dispatch (matching quicksort.hpp). Non-x86 now calls Highway VQSelect before falling back to scalar introselect. Adds HWY_QSELECT_MIN_N=16384 threshold — below this VQSelect dispatch overhead dominates on L1-resident arrays. Removes the dead DISABLE_HIGHWAY_OPTIMIZATION macro (defined but never referenced). - meson.build: register the new dispatch file for ASIMD and VSX2 targets. np.argpartition is not covered; VQSelect is keys-only and cannot maintain a parallel index array. Benchmark: M1 Pro, Apple Silicon NEON, numpy 2.4.4 scalar vs this branch dtype n k scalar vqselect speedup ------------------------------------------------- float32 100K 10 5.0 2.3 ns/e 2.2x float32 1M 100 6.8 2.6 ns/e 2.6x float32 10M 100 8.9 2.7 ns/e 3.3x float64 100K 10 8.0 3.6 ns/e 2.2x float64 1M 100 9.2 4.5 ns/e 2.0x float64 10M 100 9.3 4.5 ns/e 2.1x int32 100K 10 6.3 1.3 ns/e 4.9x int32 1M 100 5.8 1.5 ns/e 3.8x int64 1M 100 5.9 2.9 ns/e 2.0x ------------------------------------------------- n < 16K: threshold active, falls back to scalar, parity
np.partition on aarch64 has always fallen through to a 1998 scalar quickselect. np.sort already uses Highway VQSort (PR numpy#24018); this patch brings the same acceleration to np.partition via VQSelectStatic. On non-x86 platforms the quickselect_dispatch function now tries Highway VQSelect (covering int32/uint32/int64/uint64/float32/float64) before falling back to scalar introselect, matching the existing x86/non-x86 split in quicksort.hpp. A new dispatch file mirrors highway_qsort.dispatch.cpp and is registered for ASIMD and VSX2 targets in meson.build. A minimum-size threshold (HWY_QSELECT_MIN_N = 16384) skips VQSelect on L1-resident arrays where dispatch overhead outweighs the vectorisation gain; the dead DISABLE_HIGHWAY_OPTIMIZATION macro is removed. np.argpartition is not covered; VQSelect is keys-only and cannot maintain a parallel index array. Benchmark: M1 Pro NEON, numpy 2.4.4 scalar vs this branch (ns/elem) dtype n k scalar vqselect speedup ------------------------------------------------- float32 100K 10 5.0 2.3 2.2x float32 1M 100 6.8 2.6 2.6x float32 10M 100 8.9 2.7 3.3x float64 100K 10 8.0 3.6 2.2x float64 1M 100 9.2 4.5 2.0x float64 10M 100 9.3 4.5 2.1x int32 100K 10 6.3 1.3 4.9x int32 1M 100 5.8 1.5 3.8x int64 1M 100 5.9 2.9 2.0x ------------------------------------------------- n < 16K: threshold active, falls back to scalar, parity
np.partition on aarch64 has always fallen through to a 1998 scalar quickselect. np.sort already uses Highway VQSort (PR numpy#24018); this patch brings the same acceleration to np.partition via VQSelectStatic. On non-x86 platforms the quickselect_dispatch function now tries Highway VQSelect (covering int32/uint32/int64/uint64/float32/float64) before falling back to scalar introselect, matching the existing x86/non-x86 split in quicksort.hpp. A new dispatch file mirrors highway_qsort.dispatch.cpp and is registered for ASIMD and VSX2 targets in meson.build. A minimum-size threshold (HWY_QSELECT_MIN_N = 16384) skips VQSelect on L1-resident arrays where dispatch overhead outweighs the vectorisation gain; the dead DISABLE_HIGHWAY_OPTIMIZATION macro is removed. np.argpartition is not covered; VQSelect is keys-only and cannot maintain a parallel index array. Benchmark: M1 Pro NEON, numpy 2.4.4 scalar vs this branch (ns/elem) dtype n k scalar vqselect speedup ------------------------------------------------- float32 100K 10 5.0 2.3 2.2x float32 1M 100 6.8 2.6 2.6x float32 10M 100 8.9 2.7 3.3x float64 100K 10 8.0 3.6 2.2x float64 1M 100 9.2 4.5 2.0x float64 10M 100 9.3 4.5 2.1x int32 100K 10 6.3 1.3 4.9x int32 1M 100 5.8 1.5 3.8x int64 1M 100 5.9 2.9 2.0x ------------------------------------------------- n < 16K: threshold active, falls back to scalar, parity
np.partition on aarch64 has always fallen through to a 1998 scalar quickselect. np.sort already uses Highway VQSort (PR numpy#24018); this patch brings the same acceleration to np.partition via VQSelectStatic. On non-x86 platforms the quickselect_dispatch function now tries Highway VQSelect (covering int32/uint32/int64/uint64/float32/float64) before falling back to scalar introselect, matching the existing x86/non-x86 split in quicksort.hpp. A new dispatch file mirrors highway_qsort.dispatch.cpp and is registered for ASIMD and VSX2 targets in meson.build. A minimum-size threshold (HWY_QSELECT_MIN_N = 16384) skips VQSelect on L1-resident arrays where dispatch overhead outweighs the vectorisation gain; the dead DISABLE_HIGHWAY_OPTIMIZATION macro is removed. np.argpartition is not covered; VQSelect is keys-only and cannot maintain a parallel index array. Benchmark: M1 Pro NEON, numpy 2.4.4 scalar vs this branch (ns/elem) dtype n k scalar vqselect speedup ------------------------------------------------- float32 100K 10 5.0 2.3 2.2x float32 1M 100 6.8 2.6 2.6x float32 10M 100 8.9 2.7 3.3x float64 100K 10 8.0 3.6 2.2x float64 1M 100 9.2 4.5 2.0x float64 10M 100 9.3 4.5 2.1x int32 100K 10 6.3 1.3 4.9x int32 1M 100 5.8 1.5 3.8x int64 1M 100 5.9 2.9 2.0x ------------------------------------------------- n < 16K: threshold active, falls back to scalar, parity
np.partition on aarch64 has always fallen through to a 1998 scalar quickselect. np.sort already uses Highway VQSort (PR numpy#24018); this patch brings the same acceleration to np.partition via VQSelectStatic. On non-x86 platforms the quickselect_dispatch function now tries Highway VQSelect (covering int32/uint32/int64/uint64/float32/float64) before falling back to scalar introselect, matching the existing x86/non-x86 split in quicksort.hpp. A new dispatch file mirrors highway_qsort.dispatch.cpp and is registered for ASIMD and VSX2 targets in meson.build. A minimum-size threshold (HWY_QSELECT_MIN_N = 16384) skips VQSelect on L1-resident arrays where dispatch overhead outweighs the vectorisation gain; the dead DISABLE_HIGHWAY_OPTIMIZATION macro is removed. np.argpartition is not covered; VQSelect is keys-only and cannot maintain a parallel index array. Benchmark: M1 Pro NEON, numpy 2.4.4 scalar vs this branch (ns/elem) dtype n k scalar vqselect speedup ------------------------------------------------- float32 100K 10 5.0 2.3 2.2x float32 1M 100 6.8 2.6 2.6x float32 10M 100 8.9 2.7 3.3x float64 100K 10 8.0 3.6 2.2x float64 1M 100 9.2 4.5 2.0x float64 10M 100 9.3 4.5 2.1x int32 100K 10 6.3 1.3 4.9x int32 1M 100 5.8 1.5 3.8x int64 1M 100 5.9 2.9 2.0x ------------------------------------------------- n < 16K: threshold active, falls back to scalar, parity
np.partition on aarch64 has always fallen through to a 1998 scalar quickselect. np.sort already uses Highway VQSort (PR numpy#24018); this patch brings the same acceleration to np.partition via VQSelectStatic. On non-x86 platforms the quickselect_dispatch function now tries Highway VQSelect (covering int32/uint32/int64/uint64/float32/float64) before falling back to scalar introselect, matching the existing x86/non-x86 split in quicksort.hpp. A new dispatch file mirrors highway_qsort.dispatch.cpp and is registered for ASIMD and VSX2 targets in meson.build. A minimum-size threshold (HWY_QSELECT_MIN_N = 16384) skips VQSelect on L1-resident arrays where dispatch overhead outweighs the vectorisation gain; the dead DISABLE_HIGHWAY_OPTIMIZATION macro is removed. np.argpartition is not covered; VQSelect is keys-only and cannot maintain a parallel index array. Benchmark: M1 Pro NEON, numpy 2.4.4 scalar vs this branch (ns/elem) dtype n k scalar vqselect speedup ------------------------------------------------- float32 100K 10 5.0 2.3 2.2x float32 1M 100 6.8 2.6 2.6x float32 10M 100 8.9 2.7 3.3x float64 100K 10 8.0 3.6 2.2x float64 1M 100 9.2 4.5 2.0x float64 10M 100 9.3 4.5 2.1x int32 100K 10 6.3 1.3 4.9x int32 1M 100 5.8 1.5 3.8x int64 1M 100 5.9 2.9 2.0x ------------------------------------------------- n < 16K: threshold active, falls back to scalar, parity
np.partition on aarch64 has always fallen through to a 1998 scalar quickselect. np.sort already uses Highway VQSort (PR numpy#24018); this patch brings the same acceleration to np.partition via VQSelectStatic. On non-x86 platforms the quickselect_dispatch function now tries Highway VQSelect (covering int32/uint32/int64/uint64/float32/float64) before falling back to scalar introselect, matching the existing x86/non-x86 split in quicksort.hpp. A new dispatch file mirrors highway_qsort.dispatch.cpp and is registered for ASIMD and VSX2 targets in meson.build. A minimum-size threshold (HWY_QSELECT_MIN_N = 16384) skips VQSelect on L1-resident arrays where dispatch overhead outweighs the vectorisation gain; the dead DISABLE_HIGHWAY_OPTIMIZATION macro is removed. np.argpartition is not covered; VQSelect is keys-only and cannot maintain a parallel index array. Benchmark: M1 Pro NEON, numpy 2.4.4 scalar vs this branch (ns/elem) dtype n k scalar vqselect speedup ------------------------------------------------- float32 100K 10 5.0 2.3 2.2x float32 1M 100 6.8 2.6 2.6x float32 10M 100 8.9 2.7 3.3x float64 100K 10 8.0 3.6 2.2x float64 1M 100 9.2 4.5 2.0x float64 10M 100 9.3 4.5 2.1x int32 100K 10 6.3 1.3 4.9x int32 1M 100 5.8 1.5 3.8x int64 1M 100 5.9 2.9 2.0x ------------------------------------------------- n < 16K: threshold active, falls back to scalar, parity
np.partition on aarch64 has always fallen through to a 1998 scalar quickselect. np.sort already uses Highway VQSort (PR numpy#24018); this patch brings the same acceleration to np.partition via VQSelectStatic. On non-x86 platforms the quickselect_dispatch function now tries Highway VQSelect (covering int16/uint16/int32/uint32/int64/uint64/ float32/float64) before falling back to scalar introselect, matching the existing x86/non-x86 split in quicksort.hpp. A single dispatch file mirrors highway_qsort.dispatch.cpp and is registered for ASIMD and VSX2 targets in meson.build; 16-bit integers ride the same file since they are plain integers covered by base ASIMD (no ASIMDHP target needed). float16 is excluded -- HWY_HAVE_FLOAT16 gating adds complexity and float16 partition is uncommon; scalar introselect is correct there. A minimum-size threshold (HWY_QSELECT_MIN_N = 16384) skips VQSelect on L1-resident arrays where dispatch overhead outweighs the vectorisation gain; the dead DISABLE_HIGHWAY_OPTIMIZATION macro is removed. np.argpartition is not covered; VQSelect is keys-only and cannot maintain a parallel index array. Threshold rationale: a size sweep on M1 Pro (N=64-65536, kth=N/2) showed scalar introselect's random-pivot variance dominates below 16384 -- at N=8192 int64 regresses to 0.28x and float32 to 0.70x. At N=16384 every type reliably wins (1.2x for int64, up to 11x for 16-bit types), so a single threshold is used for all widths. Benchmark: M1 Pro NEON, numpy 2.4.4 scalar vs this branch (ns/elem) dtype n k scalar vqselect speedup ------------------------------------------------- float32 100K 10 5.0 2.3 2.2x float32 1M 100 6.8 2.6 2.6x float32 10M 100 8.9 2.7 3.3x float64 100K 10 8.0 3.6 2.2x float64 1M 100 9.2 4.5 2.0x float64 10M 100 9.3 4.5 2.1x int32 100K 10 6.3 1.3 4.9x int32 1M 100 5.8 1.5 3.8x int64 1M 100 5.9 2.9 2.0x ------------------------------------------------- n < 16K: threshold active, falls back to scalar, parity
np.partition on aarch64 has always fallen through to a 1998 scalar quickselect. np.sort already uses Highway VQSort (PR numpy#24018); this patch brings the same acceleration to np.partition via VQSelectStatic. On non-x86 platforms the quickselect_dispatch function now tries Highway VQSelect (covering int16/uint16/int32/uint32/int64/uint64/ float32/float64) before falling back to scalar introselect, matching the existing x86/non-x86 split in quicksort.hpp. A single dispatch file mirrors highway_qsort.dispatch.cpp and is registered for ASIMD and VSX2 targets in meson.build; 16-bit integers ride the same file since they are plain integers covered by base ASIMD (no ASIMDHP target needed). float16 is excluded -- HWY_HAVE_FLOAT16 gating adds complexity and float16 partition is uncommon; scalar introselect is correct there. A minimum-size threshold (HWY_QSELECT_MIN_N = 16384) skips VQSelect on L1-resident arrays where dispatch overhead outweighs the vectorisation gain; the dead DISABLE_HIGHWAY_OPTIMIZATION macro is removed. np.argpartition is not covered; VQSelect is keys-only and cannot maintain a parallel index array. Threshold rationale: a size sweep on M1 Pro (N=64-65536, kth=N/2) showed scalar introselect's random-pivot variance dominates below 16384 -- at N=8192 int64 regresses to 0.28x and float32 to 0.70x. At N=16384 every type reliably wins (1.2x for int64, up to 11x for 16-bit types), so a single threshold is used for all widths. Benchmark: M1 Pro NEON, numpy 2.4.4 scalar vs this branch (ns/elem) dtype n k scalar vqselect speedup ------------------------------------------------- float32 100K 10 5.0 2.3 2.2x float32 1M 100 6.8 2.6 2.6x float32 10M 100 8.9 2.7 3.3x float64 100K 10 8.0 3.6 2.2x float64 1M 100 9.2 4.5 2.0x float64 10M 100 9.3 4.5 2.1x int32 100K 10 6.3 1.3 4.9x int32 1M 100 5.8 1.5 3.8x int64 1M 100 5.9 2.9 2.0x ------------------------------------------------- n < 16K: threshold active, falls back to scalar, parity
np.sort already uses Highway VQSort (PR numpy#24018); this patch brings the same acceleration to np.partition via VQSelectStatic. On non-x86 platforms the quickselect_dispatch function now tries Highway VQSelect (covering int16/uint16/int32/uint32/int64/uint64/ float32/float64) before falling back to scalar introselect, matching the existing x86/non-x86 split in quicksort.hpp. A single dispatch file mirrors highway_qsort.dispatch.cpp and is registered for ASIMD and VSX2 targets in meson.build; 16-bit integers ride the same file since they are plain integers covered by base ASIMD (no ASIMDHP target needed). float16 is excluded -- HWY_HAVE_FLOAT16 gating adds complexity and float16 partition is uncommon; scalar introselect is correct there. A minimum-size threshold (HWY_QSELECT_MIN_N = 16384) skips VQSelect on L1-resident arrays where dispatch overhead outweighs the vectorisation gain; the dead DISABLE_HIGHWAY_OPTIMIZATION macro is removed. np.argpartition is not covered; VQSelect is keys-only and cannot maintain a parallel index array. Threshold rationale: a size sweep on M1 Pro (N=64-65536, kth=N/2) showed scalar introselect's random-pivot variance dominates below 16384 -- at N=8192 int64 regresses to 0.28x and float32 to 0.70x. At N=16384 every type reliably wins (1.2x for int64, up to 11x for 16-bit types), so a single threshold is used for all widths. Benchmark: M1 Pro NEON, numpy 2.4.4 scalar vs this branch (ns/elem) dtype n k scalar vqselect speedup ------------------------------------------------- float32 100K 10 5.0 2.3 2.2x float32 1M 100 6.8 2.6 2.6x float32 10M 100 8.9 2.7 3.3x float64 100K 10 8.0 3.6 2.2x float64 1M 100 9.2 4.5 2.0x float64 10M 100 9.3 4.5 2.1x int32 100K 10 6.3 1.3 4.9x int32 1M 100 5.8 1.5 3.8x int64 1M 100 5.9 2.9 2.0x ------------------------------------------------- n < 16K: threshold active, falls back to scalar, parity
np.sort already uses Highway VQSort (PR numpy#24018); this patch brings the same acceleration to np.partition via VQSelectStatic. On non-x86 platforms the quickselect_dispatch function now tries Highway VQSelect (covering int16/uint16/int32/uint32/int64/uint64/ float32/float64) before falling back to scalar introselect, matching the existing x86/non-x86 split in quicksort.hpp. A single dispatch file mirrors highway_qsort.dispatch.cpp and is registered for ASIMD and VSX2 targets in meson.build; 16-bit integers ride the same file since they are plain integers covered by base ASIMD (no ASIMDHP target needed). float16 is excluded -- HWY_HAVE_FLOAT16 gating adds complexity and float16 partition is uncommon; scalar introselect is correct there. A minimum-size threshold (HWY_QSELECT_MIN_N = 16384) skips VQSelect on L1-resident arrays where dispatch overhead outweighs the vectorisation gain; the dead DISABLE_HIGHWAY_OPTIMIZATION macro is removed. np.argpartition is not covered; VQSelect is keys-only and cannot maintain a parallel index array. Threshold rationale: a size sweep on M1 Pro (N=64-65536, kth=N/2) showed scalar introselect's random-pivot variance dominates below 16384 -- at N=8192 int64 regresses to 0.28x and float32 to 0.70x. At N=16384 every type reliably wins (1.2x for int64, up to 11x for 16-bit types), so a single threshold is used for all widths. Benchmark: M1 Pro NEON, numpy 2.4.4 scalar vs this branch (ns/elem) dtype n k scalar vqselect speedup ------------------------------------------------- float32 100K 10 5.0 2.3 2.2x float32 1M 100 6.8 2.6 2.6x float32 10M 100 8.9 2.7 3.3x float64 100K 10 8.0 3.6 2.2x float64 1M 100 9.2 4.5 2.0x float64 10M 100 9.3 4.5 2.1x int32 100K 10 6.3 1.3 4.9x int32 1M 100 5.8 1.5 3.8x int64 1M 100 5.9 2.9 2.0x ------------------------------------------------- n < 16K: threshold active, falls back to scalar, parity
np.sort already uses Highway VQSort (PR numpy#24018); this patch brings the same acceleration to np.partition via VQSelectStatic. On non-x86 platforms the quickselect_dispatch function now tries Highway VQSelect (covering int16/uint16/int32/uint32/int64/uint64/ float32/float64) before falling back to scalar introselect, matching the existing x86/non-x86 split in quicksort.hpp. A single dispatch file mirrors highway_qsort.dispatch.cpp and is registered for ASIMD and VSX2 targets in meson.build; 16-bit integers ride the same file since they are plain integers covered by base ASIMD (no ASIMDHP target needed). float16 is excluded -- HWY_HAVE_FLOAT16 gating adds complexity and float16 partition is uncommon; scalar introselect is correct there. A minimum-size threshold (HWY_QSELECT_MIN_N = 16384) skips VQSelect on L1-resident arrays where dispatch overhead outweighs the vectorisation gain; the dead DISABLE_HIGHWAY_OPTIMIZATION macro is removed. np.argpartition is not covered; VQSelect is keys-only and cannot maintain a parallel index array. Threshold rationale: a size sweep on M1 Pro (N=64-65536, kth=N/2) showed scalar introselect's random-pivot variance dominates below 16384 -- at N=8192 int64 regresses to 0.28x and float32 to 0.70x. At N=16384 every type reliably wins (1.2x for int64, up to 11x for 16-bit types), so a single threshold is used for all widths. Benchmark: M1 Pro NEON, numpy 2.4.4 scalar vs this branch (ns/elem) dtype n k scalar vqselect speedup ------------------------------------------------- float32 100K 10 5.0 2.3 2.2x float32 1M 100 6.8 2.6 2.6x float32 10M 100 8.9 2.7 3.3x float64 100K 10 8.0 3.6 2.2x float64 1M 100 9.2 4.5 2.0x float64 10M 100 9.3 4.5 2.1x int32 100K 10 6.3 1.3 4.9x int32 1M 100 5.8 1.5 3.8x int64 1M 100 5.9 2.9 2.0x ------------------------------------------------- n < 16K: threshold active, falls back to scalar, parity
np.sort already uses Highway VQSort (PR numpy#24018); this patch brings the same acceleration to np.partition via VQSelectStatic. On non-x86 platforms the quickselect_dispatch function now tries Highway VQSelect (covering int16/uint16/int32/uint32/int64/uint64/ float32/float64) before falling back to scalar introselect, matching the existing x86/non-x86 split in quicksort.hpp. A single dispatch file mirrors highway_qsort.dispatch.cpp and is registered for ASIMD and VSX2 targets in meson.build; 16-bit integers ride the same file since they are plain integers covered by base ASIMD (no ASIMDHP target needed). float16 is excluded -- HWY_HAVE_FLOAT16 gating adds complexity and float16 partition is uncommon; scalar introselect is correct there. A minimum-size threshold (HWY_QSELECT_MIN_N = 16384) skips VQSelect on L1-resident arrays where dispatch overhead outweighs the vectorisation gain; the dead DISABLE_HIGHWAY_OPTIMIZATION macro is removed. np.argpartition is not covered; VQSelect is keys-only and cannot maintain a parallel index array. Threshold rationale: a size sweep on M1 Pro (N=64-65536, kth=N/2) showed scalar introselect's random-pivot variance dominates below 16384 -- at N=8192 int64 regresses to 0.28x and float32 to 0.70x. At N=16384 every type reliably wins (1.2x for int64, up to 11x for 16-bit types), so a single threshold is used for all widths. Benchmark: M1 Pro NEON, numpy 2.4.4 scalar vs this branch (ns/elem) dtype n k scalar vqselect speedup ------------------------------------------------- float32 100K 10 5.0 2.3 2.2x float32 1M 100 6.8 2.6 2.6x float32 10M 100 8.9 2.7 3.3x float64 100K 10 8.0 3.6 2.2x float64 1M 100 9.2 4.5 2.0x float64 10M 100 9.3 4.5 2.1x int32 100K 10 6.3 1.3 4.9x int32 1M 100 5.8 1.5 3.8x int64 1M 100 5.9 2.9 2.0x ------------------------------------------------- n < 16K: threshold active, falls back to scalar, parity
Introduces Highway, a SIMD library from Google. Highway has it's own dispatch mechanism and SIMD flavour (which adds padding) which may be a good replacement for NumPy intrinsics. Highway also has its own set of vectorised Math routines we could bolt on.
For this patch, I've integrated the VQSort algorithm only on AArch64 so as not to impact other architectures. The VQSort algorithms performance is comparable to the AVX512 algorithms, according to this report: https://github.com/Voultapher/sort-research-rs/blob/main/writeup/intel_avx512/text.md
Performance improvements in the NumPy benchmark suite are as follows: