@MaanasArora Ping.

@dmatth1 Could you please run bench_function_base.Partition.time_partition benchmarks in the numpy repository and report the changes on you mac?

Hey, I ran bench_function_base.Partition.time_partition benchmarks on my mac:

Ratio = aarch64 with VQSelect / scalar baseline

Current ungated Highway VQSelect

Good wins on random but the big regressions come from sorted-block and ordered. We can address these by introducing a gate to fall through to scalar when near-sorted or ordered:

Gated VQSelect

The gate introduces a bit of custom logic on the non-x86 path but nothing crazy. Do you prefer to add the gate addressing the ordered and sorted-block regressions? Or are those usecases pretty rare

Thanks, the benchmarks show clear boosts on your platform! @r-devulap has a better idea about the specialized libraries involved, so I'd defer mostly to him, but including VQSelect makes sense to me as long as the performance is consistent across all non-x86 platforms...

The gate introduces a bit of custom logic on the non-x86 path but nothing crazy. Do you prefer to add the gate addressing the ordered and sorted-block regressions? Or are those usecases pretty rare

The second benchmarks have been run with the gate active, correct? Is this essentially be a single-pass check over the array? I'm not sure if the overhead (and complexity) would be worth it, especially for large arrays (there seems to be a 10-20% overhead when the gate is run), but I may be wrong. Though sorting ordered arrays seems a less common use case at least.1

Footnotes

1. Though near-sorted arrays are perhaps more common than sorted, and I'm not sure if the benchmarks would fare similarly poorly there. Probably worth extending the benchmarks on that front. Either way, probably not a prohibitive concern. ↩

Correct the gate I'd propose is a single pass over the array counting the number of descents and only calling VQSelect when not ordered or sorted-block for floats. Does add that 10-20% overhead you called out so definitely a question of whether it's worth it for the added complexity and specificity. Something like

  // one descent-count pass drives both guards
  npy_intp descents = 0;
  for (npy_intp i = 1; i < num; i++) descents += (v[i - 1] > v[i]);                                                                                                                                                           

  // Guard 1 — already-partitioned fast-path
  if (descents == 0) {                       // non-decreasing                                                                                                                                                                
      bool no_op = true;
      if constexpr (std::is_floating_point_v<T>) {   // NaN compares false -> re-check
          npy_intp nan_count = 0;
          for (npy_intp i = 0; i < num; i++) nan_count += (v[i] != v[i]);
          no_op = (nan_count == 0);
      }
      if (no_op) return true;                 // already partitioned: postcondition met, no work
  }

  // Guard 2 — dtype-aware disorder gate
  bool use_highway = true;
  if constexpr (std::is_floating_point_v<T> || sizeof(T) == 8) {                                                                                                                                                              
      use_highway = (descents > num / 4);     // float/64-bit: scalar introselect wins near-sorted
  }                                           // 16/32-bit ints skip the gate (VQSelect wins anyway)
  if (use_highway) {
      #include "highway_qselect.dispatch.h"
      NPY_CPU_DISPATCH_CALL_XB(dispfunc = np::highway::qsort_simd::template QSelect, <TF>);
  }

Regarding other non-x86 platforms I think the main concern I would have there are some of the constants I've defined, HWY_QSELECT_MIN_N = 16384 and descents > num/4 - not entirely sure that they would hold for Graviton or ppc64/VSX2 (what's covered in the Highway qselect dispatch for numpy).

Though near-sorted arrays are perhaps more common than sorted, and I'm not sure if the benchmarks would fare similarly poorly there. Probably worth extending the benchmarks on that front. - Agreed!

I am wondering if this speed difference isn't something to report to highway and maybe fix there?
The gate seems pretty simple, but it's still annoying that we would do that in NumPy.

I am wondering if this speed difference isn't something to report to highway and maybe fix there?

+1. I think highway should ideally handle these regressions and fix them inside the library rather than bringing that complexity to NumPy. Perhaps we can explore that option first?

I believe this problem of VQSelect handling near-sorted and sorted poorly is a function of Neon vs x86 SIMD differences (2 lanes vs 8) so I'm not sure how Highway would address it beyond adding the gate I recommended - and that might have a negative downstream impact on other consumers of the library outside of Numpy or those with different data needs. Also makes us wait a little bit for Numpy to vendor the new version of Highway. However if we think Highway should be the ones to address the perf gap I'll open an issue!

Note: While working on this PR I discovered a corruption bug with VQSelect and NaN which could prevent us from merging this PR. I just opened an issue with Highway to address it google/highway#3098 so we might need 2 PR's in Highway.

While working on this PR I discovered a corruption bug with VQSelect and NaN which could prevent us from merging this PR. I just opened an issue with Highway to address it google/highway#3098 so we might need 2 PR's in Highway.

That means we are missing a test in NumPy. It would be great if you can add that test case here.

one duplicating comment, but maybe you can open an issue about the speed difference at highway as well? Since if they remove any speed difference, then this PR is truly unblocked, I guess.

Sure I can add the test @r-devulap. I opened the issue with Highway about the speed regression (google/highway#3139), will monitor that and see where that leads us

Just a quick bump that this has a conflict now.