NVSHMEM 核心概念：对称内存、PE、Team 与操作语义
好，咱们今天来啃 NVSHMEM 最核心的几个概念。说实话，我刚接触 NVSHMEM 的时候，也被这些术语绕得有点晕。但后来我发现，只要搞懂了这四个东西——对称内存、PE、Team 和 操作语义——整个框架的骨架就搭起来了。
咱们一个一个说。
1. 对称内存：所有 PE 都能直接访问的“共享空间”
对称内存是 NVSHMEM 的基石。什么叫对称？说白了，就是每个 PE（处理单元）上都有一块大小相同、地址偏移相同的内存区域。
举个例子：
// 每个 PE 上分配 1024 个 float
float *sym_buf = (float*)nvshmem_malloc(1024 * sizeof(float));
PE 0 上的 sym_buf[0] 和 PE 1 上的 sym_buf[0]，虽然物理上在不同的 GPU 显存里，但逻辑上它们是对齐的。PE 0 可以直接通过 nvshmem_float_put 把数据写到 PE 1 的 sym_buf[0] 里。
关键点：对称内存的分配必须所有 PE 一起调用，而且大小必须一致。我见过有人只在主 PE 上分配，结果其他 PE 访问时直接段错误——嗯，这个坑我踩过。
对称内存的好处很明显：你不用操心数据在哪个 PE 上，只要知道对称地址，就能直接读写。这大大简化了分布式编程的复杂度。
2. PE（处理单元）：NVSHMEM 世界里的“节点”
PE 就是 NVSHMEM 里的基本执行单元。在 GPU 场景下，一个 PE 通常对应一块 GPU。每个 PE 有一个唯一的编号，从 0 到 npes-1。
你可以通过以下方式获取自己的编号和总 PE 数：
int my_pe = nvshmem_my_pe();
int npes = nvshmem_n_pes();
我个人习惯把 my_pe 叫做“我是谁”，npes 叫做“有多少兄弟”。这样记起来特别顺。
PE 之间通过对称内存进行通信。每个 PE 可以主动发起 PUT（写）或 GET（读）操作，目标可以是任意其他 PE。这种模式叫做单边通信——目标 PE 不需要参与，完全由发起方控制。
小技巧：在调试时，我经常用 printf("PE %d/%d: hello\n", my_pe, npes); 来确认每个 PE 是否都跑起来了。这招虽然土，但很管用。
3. Team：把 PE 分组，让通信更高效
有时候你不想让所有 PE 都参与通信。比如你只有 4 块 GPU 在做计算，另外 4 块在干别的事。这时候就需要 Team 来分组。
NVSHMEM 默认有一个 NVSHMEM_TEAM_WORLD，包含所有 PE。你也可以创建自己的 Team：
nvshmem_team_t my_team;
int pe_list[] = {0, 1, 2, 3}; // 只选前 4 个 PE
nvshmem_team_create_world(my_team, 4, pe_list);
Team 的好处是：你可以对 Team 内的 PE 做集体操作（比如 barrier、broadcast），而不会影响 Team 外的 PE。我在做多 GPU 的流水线并行时，就用 Team 把不同流水线阶段的 GPU 分开，互不干扰。
注意：Team 内的 PE 编号是重新映射的。比如上面例子中，PE 0 在 Team 里的编号是 0，PE 1 是 1，以此类推。别搞混了全局 PE 和 Team 内 PE 的编号。
4. 操作语义：PUT、GET 与 AMO
NVSHMEM 提供了三种核心操作：PUT、GET 和 AMO（原子内存操作）。
PUT：把数据推给别人
PUT 操作是把本地数据写到远端 PE 的对称内存里。语法很简单：
// 把本地 src 的 10 个 float 写到 PE 1 的 dest 位置
nvshmem_float_put(dest, src, 10, 1);
PUT 是异步的。也就是说，函数返回时数据不一定已经到达远端。你需要调用 nvshmem_quiet() 来确保所有 PUT 都完成。
避坑指南：我曾经在循环里连续 PUT 数据，忘了加 nvshmem_quiet()，结果远端读到的全是旧数据。嗯，从那以后我养成了“PUT 完就 quiet”的习惯。
GET：从别人那里拉数据
GET 操作是从远端 PE 的对称内存里读取数据到本地：
// 从 PE 2 的 src 位置读 10 个 float 到本地 dest
nvshmem_float_get(dest, src, 10, 2);
GET 也是异步的。你需要调用 nvshmem_barrier_all() 或 nvshmem_quiet() 来同步。
你想想看，PUT 和 GET 的区别是什么？PUT 是“我主动给你”，GET 是“我主动去拿”。在编程模型上，PUT 更适合生产者-消费者模式，GET 更适合数据请求模式。
AMO：原子操作，无锁并发
AMO（Atomic Memory Operation）用于对远端对称内存做原子更新。常见的操作有：
nvshmem_int_atomic_add：原子加nvshmem_int_atomic_compare_swap：原子比较并交换nvshmem_int_atomic_fetch：原子读取并返回旧值
举个例子：
// 原子地把 PE 3 上的 counter 加 1
int old = nvshmem_int_atomic_add(&counter, 1, 3);
AMO 的好处是不需要锁。多个 PE 可以同时对一个变量做原子操作，不会出现数据竞争。我在实现分布式计数器时，就用 AMO 来避免加锁的开销。
我的经验：AMO 虽然方便，但频繁的原子操作会降低性能。如果可能，尽量把多次原子操作合并成一次批量操作。
知识体系总览
下面这张图展示了这四个核心概念之间的关系：
从图中可以看出：对称内存是基础，PE 是执行单元，Team 是分组机制，而 PUT/GET/AMO 则是通信手段。四者缺一不可。
总结一下
这四个概念，说白了就是：
- 对称内存：所有 PE 共享的“虚拟地址空间”
- PE：每个 GPU 就是一个 PE，有自己的编号
- Team：把 PE 分组，方便做集体操作
- PUT/GET/AMO：三种通信方式，覆盖了大部分场景
搞懂了这些，NVSHMEM 的大门就算真正打开了。下一节我们会深入代码，看看这些概念在实际的并行计算中怎么用。
一句话记住：对称内存是地，PE 是人，Team 是队伍，PUT/GET/AMO 是工具。地要平，人要齐，队伍要分好，工具要用对。
公众号:蓝海资料掘金营，微信deep3321