4. 对称内存与 PE 通信：对称内存分配、数据 put/get 操作、同步屏障

好，咱们进入正题。这一章要聊的是 NVSHMEM 最核心的东西——对称内存，以及 PE 之间怎么通过它来通信。说白了，就是解决「数据怎么放、怎么拿、怎么保证大家步调一致」这三个问题。

我刚开始接触 NVSHMEM 的时候，觉得对称内存这个概念有点绕。后来在项目里调一个多卡通信的 bug，折腾了两天才发现，其实就是对称内存的地址没对上。从那以后，我对这块的理解就深了。咱们一步步来。

4.1 对称内存：到底「对称」在哪？

先问一个问题：每个 GPU 都有自己的显存，数据是私有的。那 PE 之间怎么互相访问对方的数据？

NVSHMEM 的做法是——在启动时，每个 PE 都分配一块同样大小的内存区域。这块区域在所有 PE 上的虚拟地址是一样的。这就是「对称内存」。

举个例子：PE 0 分配了一个 1024 个 float 的数组，地址是 0x7f00...000。PE 1 也分配了同样大小的数组，地址也是 0x7f00...000。虽然物理上它们在不同的显存里，但虚拟地址是镜像的。

我个人习惯把对称内存想象成「每个 PE 手里都有一份相同的地图，地图上的坐标是统一的」。你告诉 PE 1 去坐标 (x, y) 拿东西，它就知道去哪里找。

4.2 nvshmem_malloc：分配对称内存

分配对称内存的函数很简单，就是 nvshmem_malloc。用法和 malloc 几乎一样。

// 每个 PE 都分配 1024 个 float
float *buf = (float*)nvshmem_malloc(1024 * sizeof(float));
if (buf == NULL) {
    // 处理错误
}

注意：所有 PE 都必须调用 nvshmem_malloc，而且传入的大小要一致。如果 PE 0 分配 1024 个，PE 1 分配 512 个，那就不对称了，后面通信会出问题。

释放对称内存用 nvshmem_free，也是集体操作。

4.3 数据 put/get 操作：PE 之间的快递服务

有了对称内存，PE 之间就能直接读写对方的数据了。NVSHMEM 提供了两种基本操作：put（写）和 get（读）。

你想想看，这就像是你给另一个 PE 发快递。put 是你把数据寄过去，get 是你去别人那里取数据。

4.3.1 nvshmem_float_put：把数据寄给别的 PE

// PE 0 把本地数据发送给 PE 1
float local_data[4] = {1.0, 2.0, 3.0, 4.0};
float *remote_buf;  // 对称内存指针，所有 PE 上地址相同

// PE 0 执行 put 操作
nvshmem_float_put(remote_buf, local_data, 4, 1);

参数说明：

• remote_buf：目标 PE 上的对称内存地址（注意，是虚拟地址，所有 PE 上一样）
• local_data：本地源数据
• 4：传输的元素个数
• 1：目标 PE 的编号

这个操作是单边通信。PE 0 直接往 PE 1 的内存里写数据，PE 1 不需要参与。这在 HPC 里非常有用，能省掉很多同步开销。

4.3.2 nvshmem_float_get：从别的 PE 取数据

// PE 1 从 PE 0 读取数据
float local_buf[4];
float *remote_buf;  // 对称内存指针

// PE 1 执行 get 操作
nvshmem_float_get(local_buf, remote_buf, 4, 0);

get 操作也是单边的。PE 1 直接从 PE 0 的内存里读数据，PE 0 完全不知情。

4.4 同步屏障：让大家对齐步调

单边通信虽然快，但有个问题——你怎么知道数据已经写完了？

比如 PE 0 调用了 nvshmem_float_put，数据可能还在 PCIe 总线上，或者还在 GPU 的缓存里。PE 1 如果立刻去读，可能读到的是旧数据。

这时候就需要同步屏障——nvshmem_barrier。

// 所有 PE 都执行 barrier
nvshmem_barrier_all();

nvshmem_barrier_all 会阻塞所有 PE，直到每个 PE 都到达这个调用点。说白了，就是「大家等一下，等所有人都准备好了再继续」。

它的作用有两个：

1. 保证数据可见性：barrier 之前的所有 put/get 操作都完成了，数据对目标 PE 可见。
2. 同步执行流：所有 PE 的进度对齐，避免有的 PE 跑得快、有的跑得慢。

4.5 完整示例：PE 0 发数据，PE 1 收数据

咱们把上面的知识点串起来，写一个完整的例子。

#include <nvshmem.h>
#include <stdio.h>

int main(void) {
    nvshmem_init();
    int my_pe = nvshmem_my_pe();
    int n_pes = nvshmem_n_pes();

    // 分配对称内存
    float *buf = (float*)nvshmem_malloc(4 * sizeof(float));

    if (my_pe == 0) {
        // PE 0 准备数据
        float data[4] = {10.0, 20.0, 30.0, 40.0};
        // 把数据发给 PE 1
        nvshmem_float_put(buf, data, 4, 1);
    }

    // 同步屏障：确保 put 完成
    nvshmem_barrier_all();

    if (my_pe == 1) {
        // PE 1 读取数据（此时数据已经到位）
        printf("PE 1 received: ");
        for (int i = 0; i < 4; i++) {
            printf("%.1f ", buf[i]);
        }
        printf("\n");
    }

    nvshmem_free(buf);
    nvshmem_finalize();
    return 0;
}

这个例子很简单，但包含了对称内存分配、put 操作、barrier 同步三个核心要素。你可以在实际项目中把它作为模板来用。

4.6 知识体系图

下面这张图总结了本章的核心逻辑：

这张图展示了：PE 0 通过 put 把数据写入对称内存，然后所有 PE 执行 barrier 确保数据可见，最后 PE 1 读取数据。流程清晰，没有多余的步骤。

4.7 避坑指南

最后，分享几个我实际项目中遇到的坑：

• 对称内存大小不一致：每个 PE 分配的字节数必须相同。我曾经在动态计算大小时，因为浮点精度问题导致 PE 0 和 PE 1 分配的大小差了 8 字节，结果数据错位。
• 忘记 barrier：put 之后不调 barrier，直接去读，读到的是垃圾数据。这个坑我至少踩过三次。
• 在条件分支里调 barrier：如果某个 PE 进了 if 分支调了 barrier，另一个 PE 没进，那就死锁了。barrier 必须无条件执行。

嗯，对称内存这块其实不难，但细节很多。你只要记住「地址对称、集体操作、单边通信、屏障同步」这十六个字，基本就掌握了。

---