第4章 NVSHMEM 内存模型

第四章：NVSHMEM 内存模型——对称内存的分配、释放与通信

好，咱们今天来聊聊 NVSHMEM 里最核心的概念——对称内存。说实话，我第一次接触这个名词的时候，脑子里想的是「对称？是不是左右镜像那种？」后来踩了坑才明白，这里的对称，指的是所有 GPU 上看到的内存地址是一模一样的。

你想想看，在多 GPU 编程里，最头疼的是什么？是数据在哪、怎么找到它。NVSHMEM 的对称内存，说白了就是给每个 GPU 都配了一个「镜像空间」。你在 GPU 0 上分配了一块地址 0x7f00，那么在 GPU 1、GPU 2 上，同样的地址 0x7f00 也对应着一块内存。这就好比你给每个工人发了同一把钥匙，但每个工人打开的是自己工位上的抽屉。

对称内存的分配与释放

先看怎么分配。NVSHMEM 提供了 nvshmem_malloc 这个函数，用法和标准的 malloc 几乎一样。我个人习惯在程序初始化后立即分配对称内存，因为后续的 put/get 操作都依赖它。

// 分配对称内存
void *sym_ptr = nvshmem_malloc(1024 * sizeof(float));
if (sym_ptr == NULL) {
    fprintf(stderr, "对称内存分配失败！\n");
    exit(1);
}

嗯，这里要注意：nvshmem_malloc 会在所有参与计算的 GPU 上同时分配同样大小的内存。你在 PE 0 上调用一次，PE 1、PE 2、PE 3 上都会自动分配。这就是「对称」的体现。

释放的时候用 nvshmem_free，同样是对称的——你只需要在一个 PE 上调用，所有 PE 上的内存都会被释放。

nvshmem_free(sym_ptr);

⚠️ 我曾经踩过的坑： 千万别在对称内存上调用普通的 free()！那会导致内存泄漏，甚至程序崩溃。NVSHMEM 管理的内存必须用 NVSHMEM 自己的接口释放。

对称内存的 put/get 操作

有了对称内存，我们就可以在 GPU 之间搬数据了。NVSHMEM 提供了 nvshmem_float_put 和 nvshmem_float_get 这类操作。名字很直白：put 是把数据发出去，get 是把数据拉回来。

举个例子，假设 PE 0 想把一个 float 数组发给 PE 1：

// PE 0 上的代码
float *src = (float*)nvshmem_malloc(N * sizeof(float));
float *dst = (float*)nvshmem_malloc(N * sizeof(float));

// 初始化 src 数据
for (int i = 0; i < N; i++) src[i] = i * 1.0f;

// 把 src 的数据 put 到 PE 1 的 dst 上
nvshmem_float_put(dst, src, N, 1);  // 目标 PE 是 1

这里 nvshmem_float_put 的参数依次是：目标地址、源地址、元素个数、目标 PE 编号。执行完后，PE 1 上的 dst 数组就拿到了数据。

反过来，get 操作是从别的 PE 拉数据：

// PE 0 从 PE 1 拉数据
nvshmem_float_get(local_buf, remote_buf, N, 1);

💡 我的经验： put 操作通常比 get 快一点点，因为 put 是「推」的模式，不需要等待对方响应。在带宽敏感的场景下，我优先用 put。

内存 fence 与同步

好，数据发出去了，但对方什么时候能读到？这就引出了 fence 和同步的问题。

NVSHMEM 的 put/get 操作默认是异步的。也就是说，nvshmem_float_put 返回时，数据可能还在路上，并没有真正到达目标内存。如果你马上读目标地址，可能会读到旧数据。

这时候就需要 nvshmem_fence 了。它的作用是确保之前发出的所有 put 操作都对当前 PE 可见。注意，是对当前 PE 可见，不保证其他 PE 能看到。

nvshmem_float_put(dst, src, N, 1);
nvshmem_fence();  // 确保 put 的数据已经到达 PE 1
// 现在可以安全地在 PE 0 上做其他操作了

但如果你想让 PE 1 知道数据已经准备好了，光有 fence 还不够。你需要用 nvshmem_barrier 做全局同步：

// 所有 PE 都调用 barrier
nvshmem_barrier_all();  // 等所有 PE 都到达这里才继续

barrier 是「全同步」，所有 PE 必须都执行到这行代码，才能继续往下走。这比 fence 重得多，但能保证全局一致性。

🔑 核心区别：
nvshmem_fence：轻量级，只保证当前 PE 的 put 操作完成
nvshmem_barrier：重量级，所有 PE 同步，保证全局可见

知识体系结构图

下面这张图总结了对称内存的核心逻辑，我画的时候特意把 put/get 和 fence/barrier 的关系标了出来：

实战中的避坑指南

讲几个我实际项目中遇到的坑，希望能帮你少走弯路：

对称内存不能跨进程共享：每个进程有自己的虚拟地址空间，对称内存的地址相同只是巧合。别想着把一个进程的指针传给另一个进程直接解引用。
fence 不是万能的：我曾经在一个 8 GPU 的项目里，只用 fence 做同步，结果 PE 3 读到的数据总是错的。后来换成 barrier 才解决。记住：fence 只保证当前 PE 的 put 完成，不保证其他 PE 能看到。
对称内存大小要一致：所有 PE 上分配的大小必须相同。如果你在 PE 0 上分配 1024 字节，在 PE 1 上分配 2048 字节，NVSHMEM 会报错。
释放时机要小心：确保所有 put/get 操作都完成后，再调用 nvshmem_free。我见过有人一边做 put，一边释放内存，结果数据损坏。

💡 我的小技巧： 在调试阶段，可以在每次 put 之后加一个 fence，虽然会损失一点性能，但能帮你快速定位问题。等代码稳定了，再考虑去掉多余的 fence。

好了，对称内存这块就聊到这。说白了就是三个要点：分配要对称、通信用 put/get、同步靠 fence/barrier。你只要记住这三条，NVSHMEM 的基本用法就掌握了。

📚 课程目录

第四章：NVSHMEM 内存模型——对称内存的分配、释放与通信

对称内存的分配与释放

对称内存的 put/get 操作

内存 fence 与同步

知识体系结构图

实战中的避坑指南