Argsort on int32_t severely underperforms

[KungFuJesus]

Hello, I just want to make sure I'm not doing something wrong. I stumbled upon this project when I found that IPP's radix sort had buffer size calculations that were overflowing well before the limits of an INT32_MAX and all of the "Index" variants also only used 32 bit types. To be clear, a 32 bit index variant still would be beneficial and I see there are plans to eventually implement that. It's just I need a conditional code path for when there are too many values (I have a dataset generated that has length in the billions). I noticed that when the code makes use of the AVX512 instructions for argsort, it severely under performs against the radix sort (which as far as I can tell from the assembly, is completely scalar except for an iota like scan to generate the sequential index).

The benchmarks distributed with this project do show it winning by some margins against the scalar variants in here but not by the typical order of magnitudes that were touted when this project was first introduced. Is this maybe a result of the gather data sampling mitigations in microcode? I know this heavily uses compressed store and gather, and those functions were severely limited by this.

Here's my benchmark on my hardware:

[astylinski@fedoravm builddir]$ ./benchexe --benchmark_filter=/int32_t
2023-12-03T11:30:05-05:00
Running ./benchexe
Run on (6 X 3312 MHz CPU s)
CPU Caches:
  L1 Data 32 KiB (x6)
  L1 Instruction 32 KiB (x6)
  L2 Unified 4096 KiB (x6)
  L3 Unified 16384 KiB (x1)
Load Average: 1.02, 0.68, 0.32
--------------------------------------------------------------------------------
Benchmark                                      Time             CPU   Iterations
--------------------------------------------------------------------------------
simdargsort/random_128/int32_t               959 ns          956 ns       731536
simdargsort/random_256/int32_t              2045 ns         2040 ns       343034
simdargsort/random_512/int32_t              3866 ns         3857 ns       181511
simdargsort/random_1k/int32_t               9289 ns         9267 ns        75493
simdargsort/random_5k/int32_t              50900 ns        50768 ns        13789
simdargsort/random_100k/int32_t          1764351 ns      1757047 ns          400
simdargsort/random_1m/int32_t           29256670 ns     29111150 ns           24
simdargsort/random_10m/int32_t         959237861 ns    954095660 ns            1
simdargsort/smallrange_128/int32_t           722 ns          720 ns       975799
simdargsort/smallrange_256/int32_t          2049 ns         2044 ns       342384
simdargsort/smallrange_512/int32_t          3655 ns         3646 ns       191985
simdargsort/smallrange_1k/int32_t           9169 ns         9147 ns        76497
simdargsort/smallrange_5k/int32_t          21375 ns        21320 ns        32747
simdargsort/smallrange_100k/int32_t       467703 ns       465853 ns         1502
simdargsort/smallrange_1m/int32_t        9097255 ns      9047326 ns           75
simdargsort/smallrange_10m/int32_t     183596592 ns    182645622 ns            4
simdargsort/sorted_10k/int32_t            104574 ns       104314 ns         6668
simdargsort/constant_10k/int32_t            8648 ns         8628 ns        81134
simdargsort/reverse_10k/int32_t           105977 ns       105722 ns         6558
scalarargsort/random_128/int32_t             797 ns          795 ns       882026
scalarargsort/random_256/int32_t            1736 ns         1732 ns       403901
scalarargsort/random_512/int32_t            4000 ns         3990 ns       176210
scalarargsort/random_1k/int32_t            21879 ns        21822 ns        31937
scalarargsort/random_5k/int32_t           262249 ns       261537 ns         2676
scalarargsort/random_100k/int32_t        7551632 ns      7525002 ns           93
scalarargsort/random_1m/int32_t        103352847 ns    102895746 ns            7
scalarargsort/random_10m/int32_t      1974523438 ns   1964054017 ns            1
scalarargsort/smallrange_128/int32_t         667 ns          665 ns      1051997
scalarargsort/smallrange_256/int32_t        1498 ns         1494 ns       469065
scalarargsort/smallrange_512/int32_t        3152 ns         3144 ns       223444
scalarargsort/smallrange_1k/int32_t        13730 ns        13697 ns        50943
scalarargsort/smallrange_5k/int32_t       127606 ns       127285 ns         5499
scalarargsort/smallrange_100k/int32_t    3031465 ns      3019854 ns          232
scalarargsort/smallrange_1m/int32_t     38177852 ns     37995219 ns           18
scalarargsort/smallrange_10m/int32_t   695700305 ns    692244521 ns            1
scalarargsort/sorted_10k/int32_t           92423 ns        92202 ns         7578
scalarargsort/constant_10k/int32_t         83139 ns        82931 ns         8439
scalarargsort/reverse_10k/int32_t          74214 ns        74031 ns         9359

[r-devulap]

HI @KungFuJesus. We haven't compared performance of avx512_argsort with IPP SortRadixIndexAscend yet. I will try to do this sometime next week. avx512_argsort used to leverage gather instruction heavily, but because of the DOWNFALL security vulnerability, we replaced that with a scalar emulation of it in #65 and I believe that has impacted its performance. Also worthy to note that IPP sort uses an additional O(N) space whereas avx512_argsort is O(1) space, so I am not sure if this is an apples-to-apples comparison in the first place.

Your benchmark numbers look right, and I don't believe we claimed avx512_argsort to be the orders of magnitude faster than its scalar counterpart.

[KungFuJesus]

I'm sure this project never made the claim, I just recall Phoronix touting something like a 10x gain when numpy used it in place of std::sort. I'm sure there'll be nuance there. The scalar replacement does explain why me disabling GDS mitigations didn't seem to make a difference, though. The additional memory requirement is acceptable in my case. I do wonder if the sorting network approach could be applied to a radix sort for the moves. E.g. instead of comparisons for the sorting network, the bucket positions get computed. I'm like 90% sure IPP's using purely scalar operations for everything but a linear sequence generation.

[r-devulap]

I'm sure this project never made the claim, I just recall Phoronix touting something like a 10x gain when numpy used it in place of std::sort

I believe that was for quicksort, not argsort.

The additional memory requirement is acceptable in my case.

If additional space is acceptable, then you could try using key-value sort keyvalue_qsort instead of argsort. Key-value sort will sort two vectors based on the value of one of them. Pass a copy of your array and an array of indices initialized to [0, 1, 2, .. N-1]. Unlike argsort, key-value sort doesn't rely on gather (or its emulation) to load data into registers and should, in theory, be faster. It does have the downside that you have make a copy of the array which can be expensive. Would be curious to know how it performs.

[KungFuJesus]

I had tried that initially to try to get 32 bit values for indices and it wasn't working for me correctly for whatever reason. I didn't try it with size_t indices, though.

[r-devulap]

32-bit indices weren't supported until recently: #108. Did you try it after this was merged?

[KungFuJesus]

Yes, last commit I was trying against was:
d9c9737

And it was wrong, and not wrong because it was unstable, but something very wonky. I'm computing a voxelization index and scanning through these voxel indices post sort with std::upper_bound to get their groupings. I went from something like 40k unique indices down to 2 or 3. Let me reconstruct the code I used to test, it should be a simple thing to verify correctness/incorrectness for if there was a mistake on my end.

edit: this is what I used

std::vector<uint32_t> argSort32(int32_t *vals, size_t numVals)
{
    std::vector<uint32_t> idxs(numVals);
    std::iota(idxs.begin(), idxs.end(), 0);
    x86simdsort::keyvalue_qsort<uint32_t, int32_t>(idxs.data(), vals, numVals);
    return idxs;
}

[KungFuJesus]

Yeah, even with signed 32 bit types for the index I'm getting 4 unique voxel nodes instead of 47198 like there should be for this particular set of arguments.

Also flipping the keys and argument values around results 3 unique nodes, so, I don't think that's it, either.

[KungFuJesus]

If you can't reproduce this issue with random data I can gladly provide you this data in ascii since it's just indices, but it's going to be about 123 million integers.

[r-devulap]

your argsort implementation looks incorrect. You want to sort idxs based on the values in vals: therefore vals is the key and idxs is the value.

std::vector<uint32_t> argSort32(int32_t *vals, size_t numVals)
{
    std::vector<uint32_t> idxs(numVals);
    std::iota(idxs.begin(), idxs.end(), 0);
    x86simdsort::keyvalue_qsort(vals, idxs.data(), numVals);
    return idxs;
}

Couple of things to note: numVals needs to be < UINT32_MAX. Also, I am assuming vals is a copy of your data because the key-value sort will sort both vals and idxs arrays.

[r-devulap]

idxs should have numVals unique values (basically [0,1,2, ..., numVals-1] in a jumbled up order, else there is something seriously wrong with the code.

[KungFuJesus]

numVals == ~103 million, which is certainly smaller than the 2.5ish billion value of UINT32_MAX. Like I said before, I switched around vals and idxs in case I was confusing keys and values and still ended up with only 3 unique values. I haven't printed even a small slice of the output yet, I'm mostly just looking at what my std::upper_bound scan is concluding.

Did you want a gzip compressed text document of the vals array?

[KungFuJesus]

Also, I am assuming vals is a copy of your data because the key-value sort will sort both vals and idxs arrays.

In this circumstance it's ok to be destructive to vals as this is a temporary voxel node index calculation that doesn't survive past this. It's merely there to get row indices for a bunch of 3d points to arrange spatially in a grid.

Err wait...no that's not strictly true. OK, so I may need to make an explicit copy. Let me try that.

[KungFuJesus]

Ok, yeah that was it. Thanks for the tip. Here's the timing:
radix sort time: 3825.987119, kvSort time = 4555.936716

So radix sort is still winning there by a fair margin :(. Now, I did revert to the commit just prior to your PR to see what the performance was like with actual gathers (and GDS mitigations disabled, of course), and at least with the benchmarks it looks ~1.8x faster, so maybe it might have been beating IPP's index radix sort if we could rely on the microcode mitigations not being loaded.

[KungFuJesus]

Another thing going for this is that the reason it wasn't working is I was doing what the code was already doing, deriving the post sorted values. So, there's some savings there too but maybe not enough to beat radix index sort alone.

[r-devulap]

benchmarks it looks ~1.8x faster, so maybe it might have been beating IPP's index radix sort

Yeah, that is unfortunate. Hopefully future hardware won't need these mitigations and we can have a version of argsort with the gather instruction.

For there is scope of performance improvement for the [32-bit, 32-bit] key value sort, which we are working on. Will keep this open till then.