设备端点对点通信：nvshmem_put()与nvshmem_get()

好，咱们今天聊聊NVSHMEM里最基础、也最常用的操作——点对点通信。说白了，就是两个GPU之间怎么直接传数据。

你想想看，在多GPU编程里，最头疼的是什么？数据搬运。以前用MPI，得经过CPU中转，绕一大圈。NVSHMEM不一样，它让GPU之间能直接对话，就像两个人面对面聊天，不用中间人传话。

nvshmem_put()：把数据推过去

nvshmem_put() 这个API，我习惯叫它“推操作”。意思很直白：我把数据推给你。

来看个最简单的例子：

// 设备端代码
__global__ void put_example(float *dest, float *src, int nelems) {
    int my_pe = nvshmem_my_pe();
    int peer_pe = 1;  // 假设我们要发给PE 1
    
    if (my_pe == 0) {
        // PE 0 把数据推给 PE 1
        nvshmem_float_put(dest, src, nelems, peer_pe);
    }
}

这里要注意几个关键点：

• 目标地址：dest 是远端PE上的地址，必须是NVSHMEM分配的对称内存
• 源地址：src 是本地PE上的地址
• 元素个数：nelems 是传输的float数量
• 目标PE：最后一个参数指定发给谁

nvshmem_get()：把数据拉过来

有推就有拉。nvshmem_get() 是“拉操作”——我去你那里拿数据。

__global__ void get_example(float *local_buf, float *remote_buf, int nelems) {
    int my_pe = nvshmem_my_pe();
    int peer_pe = 0;  // 从PE 0拉数据
    
    if (my_pe == 1) {
        // PE 1 从 PE 0 拉数据
        nvshmem_float_get(local_buf, remote_buf, nelems, peer_pe);
    }
}

嗯，这里有个小细节。put和get的参数顺序不一样：

我个人习惯记成：put是“我推给你”，get是“我从你那拿”。这样参数顺序就自然了——put先写你的地址，get先写我的地址。

阻塞与非阻塞语义

NVSHMEM提供了两套API：阻塞的和非阻塞的。区别在哪？

阻塞版本：调用返回时，数据已经到达目的地。简单粗暴，但会浪费等待时间。

// 阻塞put - 返回时数据已到达远端
nvshmem_float_put(dest, src, nelems, peer_pe);
// 到这里，数据保证在远端可见

非阻塞版本：调用立即返回，数据可能在路上。你需要显式等待完成。

// 非阻塞put - 立即返回
nvshmem_float_put_nbi(dest, src, nelems, peer_pe);
// 可以干点别的事...
// 然后等待完成
nvshmem_quiet();

在CUDA内核中调用这些API

这些API必须在CUDA内核里调用。为什么？因为NVSHMEM走的是GPU直连通道（NVLink或PCIe），不需要CPU介入。

来看一个完整的例子：

__global__ void peer_exchange(float *my_data, float *peer_data, 
                               int nelems, int peer_pe) {
    int tid = threadIdx.x + blockIdx.x * blockDim.x;
    
    if (tid < nelems) {
        // 每个线程处理一个元素
        // 非阻塞put
        nvshmem_float_put_nbi(&peer_data[tid], &my_data[tid], 1, peer_pe);
    }
    
    __syncthreads();  // 确保所有线程都发起了put
    
    if (tid == 0) {
        // 等待所有非阻塞操作完成
        nvshmem_quiet();
    }
}

这里有个关键点：nvshmem_quiet() 是全局同步点。它保证所有之前发起的非阻塞操作都完成了。但要注意，它只保证数据到达，不保证其他PE上的线程能看到。

避坑指南

我这些年用NVSHMEM，踩过的坑不少。分享几个典型的：

1. 忘记调用nvshmem_quiet()：非阻塞操作没等完成就开始用数据，结果读到的是垃圾值。这问题特别隐蔽，有时候跑100次才出1次错。
2. 线程间竞争：多个线程往同一个远端地址写，没有做同步。NVSHMEM不保证这种操作的原子性，数据会乱掉。
3. 对称内存大小不一致：每个PE上分配的对称内存大小必须一样。我曾经有个bug，PE 0分配了100个float，PE 1只分配了50个，结果写越界了。

嗯，说到这，我想起一个真实案例。有次做多GPU训练，数据分布不均匀，有的GPU数据多，有的少。我图省事，每个PE分配了不同大小的对称内存。结果一跑就崩，查了两天才发现是这个问题。后来统一分配最大容量，用多少取多少，问题就解决了。

性能小贴士

最后给几个实用建议：

• 批量传输优于小包：一次传1MB比传1000次1KB快得多。尽量合并小数据成大块再传。
• 用非阻塞API做计算通信重叠：先发起put/get，然后做本地计算，最后等quiet。这样通信和计算并行，效率最高。
• 注意NVLink拓扑：同一节点内的GPU之间走NVLink，延迟低带宽高。跨节点走InfiniBand，延迟高一些。设计通信模式时要考虑这个差异。

好了，点对点通信就聊到这。记住put是推，get是拉，非阻塞要等quiet。这些基础打牢了，后面学集合通信就轻松多了。