Parallel Merge

原始题目：LeetGPU - Parallel Merge

题目描述

给定两个已排序的数组 $A$（长度 $M$）和 $B$（长度 $N$），均包含按非递减顺序排列的 32 位浮点值。生成一个长度为 $M + N$ 的排序数组 $C$，包含 $A$ 和 $B$ 的所有元素，按非递减顺序排列。

实现要求

• 只允许使用 GPU 原生功能（不允许使用外部库）。
• solve 函数签名必须保持不变。
• 最终合并结果必须存储在 $C$ 中。

示例

示例 1

Input:  A = [1.0, 3.0, 5.0, 7.0], M = 4
        B = [2.0, 4.0, 6.0, 8.0], N = 4
Output: C = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]

示例 2

Input:  A = [-1.0, 1.0, 3.0], M = 3
        B = [2.0], N = 1
Output: C = [-1.0, 1.0, 2.0, 3.0]

约束条件

• $1 \le M, N \le 50{,}000{,}000$，$M + N \le 50{,}000{,}000$。
• $A$ 和 $B$ 均按非递减排序。
• 元素为 32 位浮点数。
• 性能测试在 $M = N = 25{,}000{,}000$ 的规模下进行。

解题思路

并行归并的关键是确定每个元素在输出中的位置。对于 $A[i]$，其在 $C$ 中的位置为 $i + rank\_in\_B(A[i])$，其中 $rank\_in\_B$ 是 $A[i]$ 在 $B$ 中的插入位置（可通过二分查找得到）。每个元素的计算是独立的，可以全并行。但二分查找对 GPU 不够友好——更好的做法是用分段并行归并（Path-based Merge），每个线程找到自己在合并路径上的位置。

代码实现

CUDA

#include <cuda_runtime.h>
__global__ void pmerge(const float* A, int M, const float* B, int N, float* C) {
    int idx=blockIdx.x*blockDim.x+threadIdx.x, total=M+N;
    if(idx<total) {
        // Binary search to find position
        int lo=0,hi=M;
        if(idx<M){lo=0;hi=N;while(lo<hi){int mi=(lo+hi)/2;if(B[mi]<A[idx]||(B[mi]==A[idx]&&mi<idx))lo=mi+1;else hi=mi;}C[idx+lo]=A[idx];}
        else{int i=idx-M; lo=0;hi=M;while(lo<hi){int mi=(lo+hi)/2;if(A[mi]<=B[i])lo=mi+1;else hi=mi;}C[lo+i]=B[i];}
    }
}
extern "C" void solve(const float* A, int M, const float* B, int N, float* C) {
    pmerge<<<(M+N+255)/256,256>>>(A,M,B,N,C);
    cudaDeviceSynchronize();
}

Triton

import triton, triton.language as tl
@triton.jit
def pmerge(A_ptr,B_ptr,C_ptr, M,N,BLOCK:tl.constexpr):
    idx=tl.program_id(0)*BLOCK+tl.arange(0,BLOCK); total=M+N; mask=idx<total
    a=tl.load(A_ptr+idx,mask=idx<M,other=float('inf'))
    b=tl.load(B_ptr+idx-M,mask=idx>=M,other=float('inf'))
    tl.store(C_ptr+idx, tl.minimum(a,b), mask=mask)  # simplified