第3章编程模型

编程模型：对称内存与PE概念、PE间点对点通信、集合操作与同步原语

好，咱们今天来聊聊NVSHMEM的编程模型。这部分内容，说白了就是NVSHMEM的骨架。你理解了它，后面写代码就会很顺手。

我个人习惯把NVSHMEM的编程模型拆成三个层次来看：内存模型、通信模型、同步模型。咱们一个一个来。

1. 对称内存与PE概念

先说说PE。PE就是Processing Element，处理单元。在NVSHMEM里，每个GPU就是一个PE。你想想看，一个节点里插了8块GPU，那就是8个PE。跨节点的话，PE的数量就是所有GPU的总和。

那对称内存是什么？说白了，就是每个PE上都有一块大小相同、地址偏移相同的内存区域。我举个例子：

// 每个PE分配 1024 个 float
float *sym_buf = (float*)nvshmem_malloc(1024 * sizeof(float));

这行代码执行后，PE 0上的sym_buf地址是0x7f...000，PE 1上的sym_buf地址也是0x7f...000。虽然物理上它们在不同的GPU显存里，但逻辑地址是一样的。这就是对称内存的核心——全局统一的虚拟地址空间。

重要概念：对称内存是NVSHMEM通信的基础。你只能对对称内存进行远程读写操作。普通malloc出来的内存不行。

我在项目中遇到过一个问题：有个同事用cudaMalloc分配了内存，然后直接传给nvshmem_put。结果程序跑起来就崩了。嗯，这就是没理解对称内存的概念。记住，只有nvshmem_malloc分配的内存才能用于PE间通信。

2. PE间点对点通信

点对点通信，就是两个PE之间直接传数据。NVSHMEM提供了两种方式：put（推）和get（拉）。

2.1 put操作

put操作，就是当前PE主动把数据写到另一个PE的对称内存里。看代码：

// PE 0 向 PE 1 发送数据
if (nvshmem_my_pe() == 0) {
    float src[4] = {1.0, 2.0, 3.0, 4.0};
    // 把src的数据写到PE 1的sym_buf里
    nvshmem_float_put(sym_buf, src, 4, 1);
}

这里要注意，put操作是非阻塞的。什么意思？就是nvshmem_float_put返回时，数据可能还没到达目标PE。你需要调用nvshmem_quiet()来确保所有put操作都完成了。

小技巧：我一般会在批量put之后调用一次nvshmem_quiet()，而不是每个put后面都调。这样可以减少同步开销。

2.2 get操作

get操作正好相反，是从远程PE把数据拉回来：

// PE 1 从 PE 0 读取数据
if (nvshmem_my_pe() == 1) {
    float dst[4];
    // 从PE 0的sym_buf读取4个float
    nvshmem_float_get(dst, sym_buf, 4, 0);
}

get操作也是非阻塞的。你需要调用nvshmem_barrier_all()或者nvshmem_quiet()来确保数据已经到达本地。

我曾经踩过一个坑：用get操作读取数据后，没等它完成就直接用了。结果读出来的全是0。嗯，从那以后我养成了习惯——get之后一定要加同步。

3. 集合操作与同步原语

点对点通信是两个PE之间的事。那如果所有PE都要参与呢？这时候就需要集合操作了。

3.1 广播（Broadcast）

广播，就是把一个PE的数据发给所有其他PE：

// PE 0 把数据广播给所有PE
if (nvshmem_my_pe() == 0) {
    for (int i = 0; i < 4; i++) {
        sym_buf[i] = (float)i;
    }
}
// 所有PE都调用这个函数
nvshmem_float_broadcast(sym_buf, sym_buf, 4, 0);

广播完成后，所有PE的sym_buf里都有{0.0, 1.0, 2.0, 3.0}。

3.2 归约（Reduce）

归约操作，是把所有PE的数据通过某种操作（比如求和、求最大值）合并成一个结果：

// 每个PE把自己的数据准备好
sym_buf[0] = (float)nvshmem_my_pe();
// 所有PE求和归约，结果存到PE 0
nvshmem_float_sum_reduce(sym_buf, sym_buf, 1, 0);

如果总共有4个PE，PE 0上的结果就是0+1+2+3=6.0。

注意：归约操作要求所有PE同时调用。如果某个PE没调，程序就会卡死。我在调试分布式训练框架时就遇到过这种问题——某个PE因为条件分支没走到归约调用，结果整个集群hang住了。

3.3 同步原语

同步原语，说白了就是让所有PE在某个点上等齐了再继续。NVSHMEM提供了两种：

屏障（Barrier）：所有PE都到达屏障点后，才能继续往下走。
信号量（Signal）：更细粒度的同步，可以指定某个PE等待另一个PE的信号。

屏障的用法很简单：

// 所有PE在这里等待
nvshmem_barrier_all();
// 确保所有PE都完成了之前的操作

信号量稍微复杂一点，但更灵活：

// PE 0 通知 PE 1
if (nvshmem_my_pe() == 0) {
    nvshmem_signal_set(&signal_var, 1, 1);
}
// PE 1 等待信号
if (nvshmem_my_pe() == 1) {
    while (nvshmem_signal_test(&signal_var) != 1) {
        // 忙等待
    }
}

我个人更推荐用信号量来做点对点的同步，因为屏障的开销比较大——所有PE都要参与，哪怕只有两个PE需要同步。

知识体系总览

下面这张图，我把NVSHMEM编程模型的核心概念串起来了：

这张图把三个层次的关系讲清楚了。底层是PE和对称内存，中间是点对点通信，上层是集合操作和同步。每一层都依赖下一层。

避坑指南

最后，我总结几个实战中容易踩的坑：

对称内存大小不一致：每个PE上nvshmem_malloc的大小必须相同。我曾经在调试时发现PE 0分配了1024个float，PE 1只分配了512个，结果put操作直接写越界了。
忘记同步：put/get之后不加同步，数据可能还没到位。这个坑我踩过不止一次。
集合操作不齐：广播、归约这些操作，所有PE必须同时调用。如果有PE因为条件分支没走到，整个集群就卡死了。
信号量初始值：信号量变量必须初始化为0。我见过有人忘了初始化，结果信号量一直处于"已触发"状态。

嗯，编程模型这部分就讲到这里。你把这些概念吃透了，后面写NVSHMEM代码就会很顺畅。

📚 课程目录

编程模型：对称内存与PE概念、PE间点对点通信、集合操作与同步原语

1. 对称内存与PE概念

2. PE间点对点通信

2.1 put操作

2.2 get操作

3. 集合操作与同步原语

3.1 广播（Broadcast）

3.2 归约（Reduce）

3.3 同步原语

知识体系总览

避坑指南