Interleave Arrays

原始题目：LeetGPU - Interleave Arrays

题目描述

编写一个 GPU 程序，将两个 32 位浮点数数组交错合并。给定两个输入数组 $A$ 和 $B$，每个长度为 $N$，生成一个长度为 $2N$ 的输出数组，元素在 $A$ 和 $B$ 之间交替排列：

$$output = [A[0], B[0], A[1], B[1], A[2], B[2], \dots, A[N-1], B[N-1]]$$

实现要求

• 不允许使用外部库。
• solve 函数签名必须保持不变。
• 最终结果必须存储在 output 数组中。

示例

示例 1

Input:  A = [1.0, 2.0, 3.0], B = [4.0, 5.0, 6.0]
Output: [1.0, 4.0, 2.0, 5.0, 3.0, 6.0]

示例 2

Input:  A = [10.0, 20.0], B = [30.0, 40.0]
Output: [10.0, 30.0, 20.0, 40.0]

约束条件

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

解题思路

数组交织是一种分散写（Scatter）操作——连续读但间隔写。每个线程负责一对元素 $(A[i], B[i])$，写入输出的 $2i$ 和 $2i+1$ 位置。float2 向量化可将两次 ST 合并为一次 ST.64。

代码实现

CUDA

#include <cuda_runtime.h>

// 解法一：基础逐元素交织
__global__ void interleave_kernel(const float* A, const float* B, float* output, int N) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < N) {
        output[i * 2]     = A[i];
        output[i * 2 + 1] = B[i];
    }
}

// 解法二：Grid-Stride Loop
__global__ void interleave_grid_stride(const float* A, const float* B, float* output, int N) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int stride = gridDim.x * blockDim.x;
    for (int i = idx; i < N; i += stride) {
        output[i * 2]     = A[i];
        output[i * 2 + 1] = B[i];
    }
}

// 解法三：float2 向量化写入（一次 ST.64 写 8 字节 = 2 个 float）
__global__ void interleave_float2(const float* A, const float* B, float* output, int N) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int stride = gridDim.x * blockDim.x;
    float2* out2 = (float2*)output;
    for (int i = idx; i < N; i += stride) {
        out2[i] = make_float2(A[i], B[i]);
    }
}

extern "C" void solve(const float* A, const float* B, float* output, int N) {
    int threadsPerBlock = 256;
    int blocksPerGrid = (N + threadsPerBlock - 1) / threadsPerBlock;
    interleave_kernel<<<blocksPerGrid, threadsPerBlock>>>(A, B, output, N);
    cudaDeviceSynchronize();
}

Triton

import triton
import triton.language as tl
import torch


@triton.jit
def interleave_kernel(
    A_ptr, B_ptr, output_ptr,
    N: tl.constexpr,
    BLOCK_SIZE: tl.constexpr,
):
    pid = tl.program_id(0)
    idx = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = idx < N

    a = tl.load(A_ptr + idx, mask=mask)
    b = tl.load(B_ptr + idx, mask=mask)

    tl.store(output_ptr + idx * 2, a, mask=mask)
    tl.store(output_ptr + idx * 2 + 1, b, mask=mask)