Sort

[ArthurBrussee]

Add Radix sorting kernels

Burn is missing GPU sorting kernels atm. This adds a basi

DeviceRadixSort

The kernels are based on the DeviceRadixSort in https://github.com/b0nes164/GPUSorting, though with some tweaks they seem to slightly outperform the b0nes version.

TODO: More performance numbers.

The OneSweep version would be even faster but isn't as portable. The code should work on any plane size as well.

TODO: Test on plane size == 8

Scan

The kernels include essentially a small prefix sum as well. One day cubek-scan would be great to have, but for now we can leave this as a kernel specialized to sorting.

Implicit indices

One feature CUB and co don't have is to sort with indices. This is usually done by sorting tuples of (K, V) where V[i] == i. We can save having to instantiate these values and instead just use the index at compile time. I guess in a way this amounts to a fusion of an arrange + sort kernel, maybe in the future with a powerful enough fusion system that's not needed anymore.

CubeCL: tracel-ai/cubecl#1170
Burn: tracel-ai/burn#4436

Disclaimer: Initial version was done with Claude

[nathanielsimard]

I would try to follow a bit more the Guide regarding the kernel architecture. Things like NUM_SAMPLES, items per threads, threads per block, etc. could be in a blueprint.

Also, similar to the reduce kernels, I would try to split kernels into components, so we can create a kernel that work without using plane instructions (good for the CPU backend), another that is at the plane level and finally one where the sort is done across a cube.

Also I would adopt the CubeCL naming scheme with Cube, Plane and Unit instead of Block, Warp and Lane.