Stream Compaction

原始题目：LeetGPU - Stream Compaction

题目描述

给定一个包含 $N$ 个 32 位浮点数的一维数组 $A$，将所有正元素（$A[i] > 0$）压缩到输出数组 out 的前部，保持其原始相对顺序。剩余位置填充 $0.0$。

流压缩是 GPU 的基础原语，广泛应用于渲染、稀疏计算和碰撞检测。

实现要求

• 只允许使用 GPU 原生功能（不允许使用外部库）。
• solve 函数签名必须保持不变。
• out 的前 $k$ 个位置必须包含 $A$ 中所有 $A[i] > 0$ 的 $k$ 个元素，保持原顺序。
• out 的 $k$ 到 $N-1$ 位置必须为 $0.0$。
• $A[i] = 0.0$ 的元素不被选中。

示例

Input:  A = [1.0, -2.0, 3.0, 0.0, -1.0, 4.0]
Output: out = [1.0, 3.0, 4.0, 0.0, 0.0, 0.0]

约束条件

• $1 \le N \le 100{,}000{,}000$。
• $-1000.0 \le A[i] \le 1000.0$。
• out 预分配 $N$ 个元素，初始化为 $0.0$。
• 性能测试在 $N = 50{,}000{,}000$ 的规模下进行。

解题思路

流压缩的标准 GPU 实现分三步：(1) 谓词判断（每个元素是否应被保留，产生 0/1 mask）；(2) 互斥前缀和（计算每个保留元素应去的位置）；(3) 散射写入（按前缀和结果将保留元素写到正确位置）。前缀和是性能瓶颈，使用 Blelloch 算法的 work-efficient scan 是最优选择。

代码实现

CUDA

#include <cuda_runtime.h>
__global__ void compact_kernel(const float* input, float* output, int N, int* counter) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < N && input[idx] > 0.0f) {
        int pos = atomicAdd(counter, 1);
        output[pos] = input[idx];
    }
}
extern "C" void solve(const float* input, float* output, int N) {
    int *d_counter; cudaMalloc(&d_counter, sizeof(int));
    cudaMemset(d_counter, 0, sizeof(int));
    compact_kernel<<<(N+255)/256, 256>>>(input, output, N, d_counter);
    cudaDeviceSynchronize(); cudaFree(d_counter);
}

Triton

import triton, triton.language as tl
@triton.jit
def compact_kernel(input_ptr, output_ptr, N: tl.constexpr, BLOCK: tl.constexpr):
    idx = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
    mask = idx < N
    x = tl.load(input_ptr + idx, mask=mask)
    keep = x > 0.0
    # Use cumsum to compute destination indices
    pos = tl.cumsum(keep.to(tl.int32), axis=0) - 1
    tl.store(output_ptr + tl.where(keep, pos, 0), x.to(tl.float32), mask=keep & (pos >= 0))