3、NVSHMEM 基础通信模型：对称内存模型、PE 概念、通信模式（P2P/Collective）、同步语义
好，咱们今天聊聊 NVSHMEM 最核心的几个概念。说实话，我刚接触 NVSHMEM 的时候，也被它那一套“对称内存”、“PE”之类的术语搞得有点晕。但后来我发现，只要把这三个东西搞明白——对称内存模型、PE 概念、还有通信模式——剩下的就都是水到渠成的事了。
3.1 对称内存模型：NVSHMEM 的基石
先说说对称内存。这玩意儿说白了，就是每个 GPU 上都有一块“长得一模一样”的内存区域。你想想看，如果每个 GPU 都有一块大小相同、起始地址也相同的内存，那多方便啊！
我个人习惯把对称内存想象成“镜像世界”。你在 GPU 0 上往地址 0x1000 写了个数据，GPU 1 上同样地址 0x1000 的地方，就能直接读到这个数据。当然，前提是你得用 NVSHMEM 的 API 来操作。
核心要点：对称内存的分配必须通过 nvshmem_malloc() 或 cudaMalloc() + nvshmem_ptr() 来获取对称地址。普通 cudaMalloc() 分配的内存，在别的 PE 眼里就是“透明”的，看不见。
我在项目中遇到过一个问题：有个同事直接用 cudaMalloc() 分配了内存，然后试图用 nvshmem_putmem() 往别的 GPU 上写数据。结果呢？数据根本没过去，程序还静悄悄地崩溃了。嗯，这里要注意：非对称内存不能用于 NVSHMEM 的通信操作。
3.2 PE 概念：你是谁？你在哪？
PE，全称是 Processing Element。在 NVSHMEM 的世界里，每个参与通信的 GPU 就是一个 PE。每个 PE 都有一个唯一的编号，从 0 到 npes-1。
说白了，PE 就是 NVSHMEM 里的“身份证”。你通过 nvshmem_my_pe() 知道自己是几号，通过 nvshmem_n_pes() 知道总共有多少个兄弟。
| API | 作用 | 我踩过的坑 |
|---|---|---|
nvshmem_my_pe() |
获取当前 PE 的编号 | 别在 kernel 里用，那是 host-only 的 |
nvshmem_n_pes() |
获取总 PE 数量 | 确保所有 PE 都调用了 nvshmem_init() |
我曾经犯过一个低级错误：在 kernel 里直接调用了 nvshmem_my_pe()，结果编译通过了，但运行时就报错。后来查文档才发现，这些查询函数只能在 host 端用。在 device 端，你得用 nvshmem_team_my_pe() 这类带 team 的版本。
3.3 通信模式：P2P 与 Collective
NVSHMEM 支持两种通信模式：点对点（P2P）和集合通信（Collective）。这两种模式各有各的适用场景，我分别说说。
3.4.1 点对点通信（P2P）
P2P 就是两个 PE 之间直接通信。一个发，一个收。典型的 API 有 nvshmem_putmem() 和 nvshmem_getmem()。
// 示例：PE 0 向 PE 1 发送数据
if (nvshmem_my_pe() == 0) {
nvshmem_putmem(dest_on_pe1, src_on_pe0, sizeof(int), 1);
}这里要注意：nvshmem_putmem() 是异步的。也就是说，函数返回时数据不一定已经到达目的地。你需要用 nvshmem_quiet() 来确保操作完成。
我的小技巧：在调试 P2P 通信时，我习惯在 put 之后立即调用 nvshmem_quiet()，然后再检查数据。虽然会损失一点性能，但能快速定位问题。
3.4.2 集合通信（Collective）
集合通信就是所有 PE 一起参与的操作。比如广播、规约、屏障等。典型的 API 有 nvshmem_barrier_all()、nvshmem_broadcastmem() 等。
我刚开始用集合通信时，总觉得它跟 MPI 的集合操作差不多。但后来发现有个关键区别：NVSHMEM 的集合操作默认是同步的。也就是说，nvshmem_barrier_all() 会阻塞直到所有 PE 都到达屏障点。
警告：千万别在集合操作里混入 P2P 操作而不加同步！我曾经在一个项目里，PE 0 先发了个 put，然后立刻调用了 nvshmem_barrier_all()。结果 PE 1 在屏障之后才收到数据，导致逻辑错乱。正确的做法是：在 put 之后加 nvshmem_quiet()，再调屏障。
3.5 同步语义：别让你的数据“裸奔”
同步是 NVSHMEM 里最容易出问题的地方。说白了，同步就是确保数据在正确的时间到达正确的地方。
NVSHMEM 提供了几种同步机制：
nvshmem_quiet()：等待所有未完成的 P2P 操作完成。这是最常用的同步方式。nvshmem_barrier_all()：全局屏障，所有 PE 都到达后才继续。nvshmem_fence()：确保之前的所有操作在之后的操作之前完成。这个我用的不多，但在某些场景下很有用。
我记得有一次，我在一个 8-GPU 的节点上跑程序，结果数据总是对不上。排查了半天，发现是某个 PE 在 nvshmem_putmem() 之后没有调用 nvshmem_quiet()，就直接去读目标内存了。嗯，这其实是个典型的“数据裸奔”问题。
避坑指南：我曾经在写一个多 GPU 的流水线程序时，为了追求极致性能，把 nvshmem_quiet() 给省了。结果数据错乱得一塌糊涂。后来我总结了一个原则：“宁可多同步一次，也不要少同步一次”。性能可以慢慢优化，但正确性是第一位的。
3.6 知识体系总览
下面这张图是我自己画的，把本章的核心概念串了起来。你可以把它当作一个“思维导图”来用。
这张图把对称内存、PE 概念、通信模式这三个核心要素串在了一起。你想想看，没有对称内存，数据就没法共享；没有 PE 编号，就不知道谁是谁；没有同步，数据就会乱套。三者缺一不可。
好了，这一章的内容就到这里。记住我说的：对称内存是基础，PE 是身份，同步是保障。把这三点刻在脑子里，后面的章节你会学得轻松很多。