3. 核心概念：PE、对称内存、对称堆、对称队列

好，咱们进入正题。这一章讲的是 NVSHMEM 的四个核心概念。说白了，它们是整个通信方案的基石。你如果不理解这几个东西，后面写代码肯定会踩坑。我自己刚接触 NVSHMEM 时，就在「对称内存」上栽过跟头，后面我会细说。

3.1 PE（Processing Element）—— 最小的执行单元

PE 是什么？你可以把它理解成一个「干活的人」。在 NVSHMEM 的世界里，每个 PE 就是一个独立的 GPU 或 CPU 进程。每个 PE 都有自己的编号，从 0 开始，一直到 n-1。

举个例子。假设你有 4 块 GPU，每块 GPU 上跑一个进程。那这 4 个进程就是 4 个 PE。它们之间通过 NVSHMEM 通信，不需要经过 CPU 中转。

我习惯把 PE 比作「工位」。每个工位上有一个工人，工人之间可以互相递工具（数据）。但递工具的方式有讲究——不能乱扔，得放到对方能拿到的地方。这个「地方」，就是对称内存。

3.2 对称内存（Symmetric Memory）—— 大家都能看到的共享空间

对称内存是 NVSHMEM 最核心的概念。它指的是：每个 PE 都有一块内存，但所有 PE 都能访问它。听起来像共享内存？不完全一样。

传统 MPI 里，你要发数据，得先打包、发送、对方接收。NVSHMEM 不一样。你只需要把数据写到对方的对称内存里，对方直接就能读到。省去了「发送-接收」的握手过程。

嗯，这里要注意：对称内存不是物理上共享的。它只是逻辑上可访问。每个 PE 的对称内存，物理上还是在自己本地。但 NVSHMEM 通过底层硬件（比如 NVIDIA 的 NVLink）让你能直接读写别人的内存。

3.3 对称堆（Symmetric Heap）—— 统一管理对称内存的池子

对称堆是什么？你可以把它想象成一个「大水池」。所有 PE 的对称内存，都是从这同一个池子里分配的。每个 PE 的堆大小一样，起始地址也一样。

为什么地址要一样？因为 NVSHMEM 的通信机制依赖地址映射。如果 PE 0 的变量 A 在地址 0x1000，那 PE 1 的变量 A 也必须在地址 0x1000。这样，PE 0 要写数据给 PE 1，只需要说「把数据写到地址 0x1000」，PE 1 就能在自己的 0x1000 处读到。

对称堆的大小，在初始化时通过 nvshmem_init 或环境变量设置。我个人建议，堆大小不要设得太小。否则运行时动态分配会失败。我见过有人设了 1GB，结果模型参数一加载就爆了。

3.4 对称队列（Symmetric Queue）—— 无锁的通信管道

对称队列是 NVSHMEM 提供的一种高级数据结构。它本质上是一个环形缓冲区，但所有 PE 都能往里面 push 数据，也能从里面 pop 数据。关键是：不需要加锁。

为什么不需要锁？因为 NVSHMEM 利用 GPU 的原子操作（atomic operations）来保证并发安全。每个 PE 通过原子指令更新队列的头尾指针，不会冲突。

对称队列适合做「生产者-消费者」模式。比如，PE 0 计算完一批数据，push 到队列里。PE 1 从队列里 pop 出来继续处理。整个过程不需要显式同步。

不过要注意：对称队列的大小是固定的。如果队列满了，push 操作会阻塞。如果队列空了，pop 操作也会阻塞。所以你要根据数据量合理设置队列长度。

3.5 它们之间的关系

这四个概念不是孤立的。它们的关系是这样的：

• PE 是执行者，负责干活。
• 对称内存 是 PE 之间共享的存储空间。
• 对称堆 是管理对称内存的分配器。
• 对称队列 是建立在对称内存上的通信数据结构。

你可以这样理解：PE 从对称堆里申请对称内存，然后用对称队列来传递数据。整个流程一气呵成。

下面这张图展示了它们之间的协作关系：

从图中可以看到：PE 0 和 PE 1 各自从对称堆里分配对称内存。虽然物理上它们在不同 GPU 上，但地址空间是对齐的。然后，它们可以在对称内存上构建对称队列，实现无锁通信。

3.6 代码示例：分配对称内存并写入

光说不练假把式。咱们看一段实际代码：

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

int main() {
    // 初始化 NVSHMEM
    nvshmem_init();
    
    int my_pe = nvshmem_my_pe();
    int n_pes = nvshmem_n_pes();
    
    // 分配对称内存
    int *sym_buf = (int*)nvshmem_malloc(1024 * sizeof(int));
    
    // 每个 PE 在自己的对称内存里写入数据
    for (int i = 0; i < 1024; i++) {
        sym_buf[i] = my_pe * 1000 + i;
    }
    
    // 同步，确保所有 PE 都写完了
    nvshmem_barrier_all();
    
    // PE 0 读取 PE 1 的数据
    if (my_pe == 0) {
        int val = nvshmem_int_g(sym_buf, 1);  // 从 PE 1 的 sym_buf[0] 读取
        printf("PE 0 读到 PE 1 的数据: %d\n", val);
    }
    
    // 释放对称内存
    nvshmem_free(sym_buf);
    
    nvshmem_finalize();
    return 0;
}

这段代码展示了最核心的操作：nvshmem_malloc 分配对称内存，nvshmem_int_g 远程读取。注意，nvshmem_barrier_all 是必须的，否则可能读到脏数据。

3.7 对称队列的使用示例

再来看对称队列的用法：

#include <nvshmem.h>
#include <nvshmemx.h>

int main() {
    nvshmem_init();
    
    int my_pe = nvshmem_my_pe();
    
    // 创建对称队列，容量 100 个元素
    nvshmemx_symmetric_queue_t queue;
    nvshmemx_symmetric_queue_init(&queue, 100);
    
    if (my_pe == 0) {
        // PE 0 生产数据
        int data = 42;
        nvshmemx_symmetric_queue_push(&queue, &data);
        printf("PE 0 push: %d\n", data);
    }
    
    nvshmem_barrier_all();
    
    if (my_pe == 1) {
        // PE 1 消费数据
        int data;
        nvshmemx_symmetric_queue_pop(&queue, &data);
        printf("PE 1 pop: %d\n", data);
    }
    
    nvshmemx_symmetric_queue_destroy(&queue);
    nvshmem_finalize();
    return 0;
}

你看，代码非常简洁。没有锁，没有条件变量。NVSHMEM 在底层帮你搞定了所有同步问题。

3.8 总结一下

这一章我们讲了四个核心概念：

• PE：执行单元，每个 GPU 进程一个。
• 对称内存：所有 PE 都能访问的共享内存空间。
• 对称堆：统一管理对称内存的分配器。
• 对称队列：基于对称内存的无锁通信队列。

我个人觉得，对称内存是理解 NVSHMEM 的关键。你只要搞懂了「地址对齐、远程直接访问」这个思想，后面的 put/get 操作、collective 通信，都会变得很自然。

嗯，下一章我们会深入 put/get 操作，看看数据到底是怎么在 PE 之间流动的。到时候我会结合一个实际的多卡训练案例来讲。

---