Sorting

原始题目：LeetGPU - Sorting

题目描述

编写一个 GPU 程序，对 32 位浮点数数组按升序排序。可自由选择任意排序算法。结果必须写回原数组。

实现要求

• 不允许使用外部库。
• solve 函数签名必须保持不变。

示例

Input:  [5.0, 2.0, 8.0, 1.0, 9.0, 4.0], N = 6
Output: [1.0, 2.0, 4.0, 5.0, 8.0, 9.0]

约束条件

• $1 \le N \le 1{,}000{,}000$。
• 性能测试在 $N = 1{,}000{,}000$ 下进行。

解题思路

GPU 排序的经典选择是基数排序（Radix Sort）——适合 32 位整数/浮点数，每次按若干位分桶计数，$O(N)$ 时间。对于浮点数，需要先将 IEEE 754 表示转换为可排序的整数编码（正数符号位翻转、负数整体翻转）。对于小 $N$（$\le 10^6$），双调排序（Bitonic Sort） 在共享内存中也非常高效，且不需要额外内存。

代码实现

CUDA

#include <cuda_runtime.h>
__global__ void bitonic_sort(float* data, int N) {
    extern __shared__ float shared[];
    int tid = threadIdx.x, i = blockIdx.x * blockDim.x + tid;
    shared[tid] = (i < N) ? data[i] : INFINITY;
    __syncthreads();
    for(int size = 2; size <= blockDim.x; size <<= 1) {
        for(int stride = size >> 1; stride > 0; stride >>= 1) {
            int idx = tid ^ stride;
            float other = shared[idx];
            bool ascend = ((tid / size) % 2 == 0);
            if((shared[tid] > other) == ascend && idx > tid) {
                float tmp = shared[tid]; shared[tid] = other; shared[idx] = tmp;
            }
            __syncthreads();
        }
    }
    if(i < N) data[i] = shared[tid];
}
extern "C" void solve(float* data, int N) {
    bitonic_sort<<<(N+255)/256, 256, 256*sizeof(float)>>>(data, N);
    cudaDeviceSynchronize();
}

Triton

import triton, triton.language as tl
@triton.jit
def sort_kernel(data_ptr, N, BLOCK:tl.constexpr):
    idx=tl.program_id(0)*BLOCK+tl.arange(0,BLOCK); mask=idx<N
    x=tl.load(data_ptr+idx,mask=mask,other=float('inf'))
    sorted,_=tl.sort(x,descending=False)
    tl.store(data_ptr+idx,sorted,mask=mask)