Nearest Neighbor

原始题目：LeetGPU - Nearest Neighbor

题目描述

编写一个 GPU 程序，对于设备上存储的 $N$ 个三维点，将 indices[i] 填充为距离 points[i] 最近的点的索引 $j \neq i$。比较平方欧氏距离即可——无需计算平方根：

$$d_{ij}^2 = (x_i - x_j)^2 + (y_i - y_j)^2 + (z_i - z_j)^2$$

实现要求

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

示例

Input:  points = [(0,0,0), (1,0,0), (5,5,5)], N = 3
Output: indices = [1, 0, 1]  (0↔1互为最近邻，2离1最近)

约束条件

• $1 \le N \le 100{,}000$。
• 坐标为 32 位浮点数，范围 $[-1000, 1000]$。
• 性能测试在 $N = 10{,}000$ 的规模下进行。

解题思路

最近邻搜索的朴素做法是 $O(N^2)$ 的逐对距离计算，对于大 $N$ 不可接受。GPU 上可以并行化：每个线程/block 处理一个查询点，扫描所有候选点。对于 $N = 100{,}000$，可使用分块策略：查询点分到不同 block，每个 block 将候选点分批加载到共享内存中比较，避免重复从全局内存读取。更高效的数据结构如 KD-Tree 或 Ball Tree 在 GPU 上实现较为复杂。

代码实现

CUDA

#include <cuda_runtime.h>
__global__ void nn_kernel(const float* pts, int* idxs, int N) {
    int i=blockIdx.x*blockDim.x+threadIdx.x;
    if(i<N) {
        float xi=pts[i*3],yi=pts[i*3+1],zi=pts[i*3+2];
        float min_d=INFINITY; int min_j=-1;
        for(int j=0;j<N;j++) if(j!=i) {
            float dx=xi-pts[j*3],dy=yi-pts[j*3+1],dz=zi-pts[j*3+2];
            float d=dx*dx+dy*dy+dz*dz;
            if(d<min_d){min_d=d;min_j=j;}
        }
        idxs[i]=min_j;
    }
}
extern "C" void solve(const float* pts, int* idxs, int N) {
    nn_kernel<<<(N+255)/256,256>>>(pts,idxs,N);
    cudaDeviceSynchronize();
}

Triton

import triton, triton.language as tl
@triton.jit
def nn_kernel(pts_ptr,idxs_ptr,N,BLOCK:tl.constexpr):
    i=tl.program_id(0)*BLOCK+tl.arange(0,BLOCK); mask=i<N
    xi=tl.load(pts_ptr+i*3,mask=mask);yi=tl.load(pts_ptr+i*3+1,mask=mask);zi=tl.load(pts_ptr+i*3+2,mask=mask)
    min_d=tl.full((BLOCK,),float('inf'),tl.float32); min_j=tl.full((BLOCK,),-1,tl.int32)
    for j in range(N):
        dx=xi-tl.load(pts_ptr+j*3);dy=yi-tl.load(pts_ptr+j*3+1);dz=zi-tl.load(pts_ptr+j*3+2)
        d=dx*dx+dy*dy+dz*dz; valid=(d<min_d)
        min_d=tl.where(valid&(tl.arange(0,BLOCK)!=j%BLOCK|j!=i),d,min_d)
    tl.store(idxs_ptr+i,min_j,mask=mask)