[icfaust]

Fix key-value sort stalling on 129-256 element arrays with a dominant minimum

Summary

keyvalue_qsort with 32-bit keys and 32-bit values on AVX-512 runs 8-20x slower than on distinct keys when one value
is most of a 129-256 element array (sub)range and that value is the range's minimum. It recurses to its max_iters
limit and finishes in the scalar heap_sort, and ends up 2x slower than std::sort. Zero-heavy data (sparse
features, counts, one-hot columns) hits this whenever a sort or one of its partitions lands in that size range.

The fix applies the existing median-is-smallest / median-is-largest / near-constant pivot rules of get_pivot_smart
to its small-array branch as well, which today returns the plain get_pivot median. One function, +27/-1.

What goes wrong

kvsort_ (src/xss-common-keyvaluesort.hpp, v7.0):

1. Ranges of at most 128 elements go to the bitonic network kvsort_n<..., 128>.
2. Larger ranges take a pivot from get_pivot_smart. For ranges of at most 4 * numVecs * numlanes elements
   (256 for zmm_vector<float>, <int32_t> and <uint32_t>, 16 lanes), it returns get_pivot: the median of
   numlanes evenly spaced samples, with no check for duplicates.
3. kvpartition_unrolled splits the range into [< pivot] and [>= pivot]; kvsort_ recurses on the left side
   only if pivot != smallest, and on the right side only if pivot != biggest.

If more than half of the samples equal the range's minimum, the median is the minimum. The left side is then empty
and the right side is the whole range again, so each level partitions the same elements and makes no progress until
max_iters = 2 * log2(n) runs out and heap_sort sorts the range one element at a time.

A dominant maximum does not stall: the median is the maximum, everything below it moves left, and the right side is
constant. Descending sorts behave like ascending ones, since kvsort_ sorts ascending and reverses.

Ranges above 256 elements already avoid this: get_pivot_smart handles median == smallest with
choosePivotMedianIsSmallest, median == largest with choosePivotMedianIsLargest, and an all-equal sample with
get_pivot_near_constant. Only the small-array branch skips them. They can still reach it through recursion:
partitioning a 300-element range with 50% zeros leaves a zero-heavy left part of 129-256 elements.

Where it applies

The small-array branch is reached for a range of n elements only when the sort's base case is below
4 * 4 * numlanes:

Sort	Vector type	Base case	Small-array branch	Reached
keyvalue_qsort, 32-bit key and 32-bit value, AVX-512	zmm_vector (16 lanes)	128	<= 256	129-256
keyvalue_qsort, 64-bit key, AVX-512	zmm_vector (8 lanes)	128	<= 128	no
keyvalue_qsort, mixed 32/64-bit key and value	half vectors (8 lanes)	128	<= 128	no
keyvalue_qsort, AVX2	8 or 4 lanes	128	<= 128 or 64	no
qsort, AVX-512 32-bit	16 lanes	512	<= 256	no
qsort, AVX-512 64-bit	8 lanes	256	<= 128	no
qsort, AVX-512 16-bit	32 lanes	512	<= 512	no
qsort, AVX2 32-bit	8 lanes	256	<= 128	no
qsort, AVX2 64-bit	4 lanes	64	<= 64	no
qsort<_Float16>, AVX512-FP16	32 lanes	128	<= 512	129-512
any qsort built with XSS_MINIMAL_NETWORK_SORT		numlanes	<= 16 x numlanes	yes

The change

src/xss-pivot-selection.hpp, get_pivot_smart: the small-array branch takes the same numlanes samples as
get_pivot, sorts them with vtype::sort_vec, and runs the same decision chain as the large-array branch below it:
all samples equal -> get_pivot_near_constant; median equals the smallest sample -> choosePivotMedianIsSmallest;
median equals the largest -> choosePivotMedianIsLargest; otherwise the median. The large-array branch, get_pivot
and get_pivot_blocks (used by kvselect_) are unchanged.

Results

One core of a Xeon Platinum 8480+ (Sapphire Rapids), GCC, -O2 -march=sapphirerapids, 20000 repetitions, nanoseconds
per sort. std::sort sorts (key, value) pairs, for reference. Keys: distinct (uniform), or 70% / 90% of the
array equal to the minimum (0) or the maximum (1000) with the rest uniform in 1..999 (bench_kv_duplicates.cpp).

key	order	n	keys	before	after	std::sort
float	asc	100	distinct	347	337	849
float	asc	100	70% min	327	323	860
float	asc	100	90% min	327	323	844
float	asc	100	70% max	328	323	934
float	asc	200	distinct	795	774	2020
float	asc	200	70% min	4278	354	2017
float	asc	200	90% min	4343	269	2196
float	asc	200	70% max	292	311	2222
float	asc	250	distinct	817	800	2430
float	asc	250	70% min	5225	516	2745
float	asc	250	90% min	5146	292	2709
float	asc	250	70% max	494	513	2878
float	asc	300	distinct	1293	1274	3183
float	asc	300	70% min	664	554	3466
float	asc	300	90% min	452	425	2899
float	asc	300	70% max	548	536	3813
float	asc	500	distinct	2615	2586	5464
float	asc	500	70% min	1098	1045	6227
float	asc	500	90% min	497	546	5868
float	asc	500	70% max	836	864	6506
float	desc	100	distinct	371	379	811
float	desc	100	70% min	378	377	896
float	desc	100	90% min	375	376	905
float	desc	100	70% max	381	379	1008
float	desc	200	distinct	895	901	1969
float	desc	200	70% min	4392	477	2009
float	desc	200	90% min	4451	380	2138
float	desc	200	70% max	418	393	2245
float	desc	250	distinct	974	981	2415
float	desc	250	70% min	5333	656	2849
float	desc	250	90% min	5357	436	2826
float	desc	250	70% max	657	653	2788
float	desc	300	distinct	1452	1452	3167
float	desc	300	70% min	849	744	3508
float	desc	300	90% min	619	611	2894
float	desc	300	70% max	727	726	3846
float	desc	500	distinct	3000	2974	5457
float	desc	500	70% min	1412	1404	6116
float	desc	500	90% min	837	843	5654
float	desc	500	70% max	1204	1158	6214
int32	asc	100	distinct	286	291	756
int32	asc	100	70% min	288	288	752
int32	asc	100	90% min	287	287	743
int32	asc	100	70% max	288	287	828
int32	asc	200	distinct	682	681	1864
int32	asc	200	70% min	4023	283	1753
int32	asc	200	90% min	4105	243	1805
int32	asc	200	70% max	247	235	1818
int32	asc	250	distinct	687	703	2304
int32	asc	250	70% min	5144	488	2548
int32	asc	250	90% min	5087	220	2171
int32	asc	250	70% max	456	486	2464
int32	asc	300	distinct	1112	1113	2892
int32	asc	300	70% min	586	510	2927
int32	asc	300	90% min	325	331	2416
int32	asc	300	70% max	502	483	3225
int32	asc	500	distinct	2371	2410	5155
int32	asc	500	70% min	993	975	5465
int32	asc	500	90% min	441	445	4715
int32	asc	500	70% max	767	761	5607

Testing

• test_kv_correctness.cpp: keyvalue_qsort against std::sort for float/int32 (with and without NaN and
  hasnan), int32/int32, uint32/uint32, float/int64, double/int32 and double/int64; n = 1..300 and every 37th size
  to 2000; five patterns (distinct, 90% minimum, 60% of one middle value, 80% of one value, three values); ascending
  and descending. It checks the sorted keys, NaN placement, and that the values are a permutation each carrying its
  own key. 0 mismatches on -march=sapphirerapids and -march=haswell, with and without the change.
• The repository's gtest suite was not run (gtest was not available on the machine).

Not fixed here

qsort<_Float16> on AVX512-FP16 also slows down with a dominant zero (n = 200-500, 70-90% zeros: 3.6-9.3 us against
0.6-1.1 us for distinct keys) and falls back to std::sort. It goes through the patched branch, but the change does
not help it, so the cause is elsewhere (possibly next_value<_Float16> in choosePivotMedianIsSmallest). Not
investigated further.

Context

Found in a oneDAL's decision forest training prototype, which sorts each feature's values in tree nodes of a few hundred rows
(float keys, int32 row indices). On scikit-learn_bench's RandomForestRegressor case for gisette (1000 x 5000, 87%
zeros), 49% of fit time was in kvsort_'s heap_sort fallback; with this change the fit took 0.64 s instead of
1.87 s, with identical test scores (e.g. this bug caused a x3 slowdown in the prototype in regression on short wide data).

Analysis and implementation was done by claude opus 5.5 (with feedback by me). I have reviewed the code myself and stand by it.

[icfaust]

Hey @r-devulap I know its been a while. I've been allowed to restart the project which I had tabled long ago working on integrating this code into oneDAL. I will be submitting a second PR which does a more formal enablement, this gets basic use of keyvalue store fixed for our use case. Let me know if you have any questions.

[r-devulap]

Good to hear from you again @icfaust and thanks for the PR. I will take a look soon and looking forward to your next PR.