Summary

This PR improves torch.sort and torch.unique performance by 15% to 50% on NVIDIA GPUs by optimizing CUDA register allocation in radix sort operations.

The key change: specialize OpaqueType<N> to use native integer types (uint8_t, uint16_t, uint32_t, uint64_t) for common sizes (1, 2, 4, 8 bytes) instead of char data[N]. This enables more efficient register allocation while preserving the template deduplication strategy.

The following table shows the speedup on various input shapes and GPUs. Sorting is performed on the last dimension, and baseline torch version is 2.9.0.

Analysis

torch.sort and torch.unique use radix_sort_pairs, which internally calls cub::DeviceRadixSort::SortPairs. Since values are only copied (never compared), we cast them to OpaqueType<sizeof(value_t)> to minimize template instantiations. For example, both int32 and float32 values map to the same OpaqueType<4>.

The Problem

The previous char data[N] implementation causes inefficient register allocation. Here is one reason I find from SASS code. For 8-byte types:

• char data[8]: Compiler may allocate 8 registers (one per byte)

• uint64_t data: Compiler allocates 2 registers (standard 64-bit handling)

This happens because the compiler doesn't recognize char[8] as a cohesive 64-bit value, treating each byte independently, which increases register pressure and reduces GPU occupancy.

From Nsight Compute, when using char data[8], the registers per thread is 166, and corresponding theoretical occupancy is 18.75%. When using native uint64_t, the registers per thread is 80, and corresponding theoretical occupancy is 37.5%.

The Solution

Specialize OpaqueType<N> for common sizes using native integer types:

// Before
template <int N> struct alignas(N) OpaqueType { char data[N]; };

// After
template <int N> struct alignas(N) OpaqueType { char data[N]; }; // fallback
template <> struct alignas(1) OpaqueType<1> { uint8_t data; };
template <> struct alignas(2) OpaqueType<2> { uint16_t data; };
template <> struct alignas(4) OpaqueType<4> { uint32_t data; };
template <> struct alignas(8) OpaqueType<8> { uint64_t data; };

This preserves the template deduplication strategy (all 8-byte types still use the same OpaqueType<8> instantiation) while enabling better register allocation.

Testing & Compatibility

Testing:

✅ Correctness tests pass for various input types (bfloat16, int32, float32, int64), shapes, and dimensions (1, 2, 3)
✅ Register usage reduction verified with NSight Compute
✅ Linter passes

Compatibility:

✅ No API/ABI changes
✅ Template instantiation count unchanged

Reference

For detailed analysis, please refere to my previous blog: Performance Optimization of torch.sort on GPU