第3章对称内存与 PE 概念

对称内存与 PE 概念：NVSHMEM 的基石

各位同学，今天我们来聊聊 NVSHMEM 里最基础、也最重要的两个概念——对称内存和 PE。说实话，我刚开始接触 NVSHMEM 的时候，也被这两个词绕得有点晕。但搞懂了它们，后面的原子操作、同步机制才能真的理解透。

什么是 PE？说白了就是“干活的人”

PE，全称 Processing Element，翻译过来叫“处理单元”。你想想看，在 GPU 编程里，我们有线程、有 block、有 grid。但在 NVSHMEM 的世界里，最小的独立执行单元就是 PE。

每个 PE 都有自己的身份编号，叫 my_pe。整个系统里一共有多少个 PE 呢？用 n_pes 就能拿到。我习惯把 PE 想象成一个工位上的工人——每个工人有编号，大家协同干活。

核心要点：PE 是 NVSHMEM 通信的基本单位。一个 PE 可以是一个 GPU，也可以是一个 CPU 线程。在纯 GPU 场景下，一个 PE 通常对应一个 GPU。

对称内存：大家都能看到的“共享黑板”

对称内存（Symmetric Memory）是 NVSHMEM 里最巧妙的设计之一。说白了，就是每个 PE 都有一块内存，而且这些内存的虚拟地址在所有 PE 上都是一样的。

为什么会这样设计？我举个例子。假设你有 4 个 GPU，每个 GPU 上分配了 1GB 的对称内存。那么 GPU 0 上这块内存的起始地址是 0x7f0000000000，GPU 1、GPU 2、GPU 3 上对应的地址也是 0x7f0000000000。这样，当你做远程内存访问时，只需要告诉对方“我要读写地址 0x7f0000000000”，对方就知道你指的是哪块内存。

嗯，这里要注意：虽然地址一样，但物理上它们是分布在不同的设备上的。千万别以为地址相同就是同一块物理内存。

我曾经踩过的坑：有一次我在代码里直接用了 malloc 分配内存，然后试图用 NVSHMEM 做远程访问。结果程序直接崩溃。后来才意识到，只有通过 nvshmem_malloc 分配的内存才是对称内存。普通 malloc 分配的内存，其他 PE 根本看不到。

my_pe 与 n_pes 的使用

这两个函数是 NVSHMEM 编程的“门把手”。几乎所有程序的开头都会用到它们。

my_pe：返回当前 PE 的编号，从 0 开始
n_pes：返回系统中总的 PE 数量

我个人的习惯是，在程序一开始就把它们存到局部变量里：

int my_pe = nvshmem_my_pe();
int n_pes = nvshmem_n_pes();
printf("Hello from PE %d of %d\n", my_pe, n_pes);

你想想看，如果没有这两个函数，你怎么知道自己是哪个 PE？怎么知道该跟谁通信？

对称内存的分配与释放

分配对称内存，必须用 NVSHMEM 提供的接口：

// 分配 1024 个 int 类型的对称内存
int *sym_mem = (int*)nvshmem_malloc(1024 * sizeof(int));

// 使用完毕后释放
nvshmem_free(sym_mem);

这里有个细节：nvshmem_malloc 会在所有 PE 上都分配同样大小的内存。也就是说，如果你在 PE 0 上调用了这个函数，PE 1、PE 2、PE 3 上也会自动分配好对应的内存。这就是“对称”的含义。

小技巧：我在项目中经常用对称内存来存储全局的计数器或者状态标志。因为每个 PE 都能直接读写其他 PE 上的对称内存，实现起来特别方便。

对称内存的访问模式

对称内存支持两种访问模式：

访问模式	说明	示例
本地访问	直接读写当前 PE 上的对称内存	`sym_mem[0] = 42;`
远程访问	通过 NVSHMEM API 读写其他 PE 上的对称内存	`nvshmem_int_put(&dest, &src, 1, target_pe);`

本地访问就跟操作普通内存一样，直接读写就行。远程访问则需要调用 NVSHMEM 提供的 put/get 操作。我刚开始学的时候，经常忘记区分这两种模式，结果本地访问写得很开心，远程访问却忘了用 API——程序跑起来数据全是错的。

一个完整的例子

来看一个完整的程序，感受一下对称内存和 PE 概念是怎么用的：

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

int main(void) {
    nvshmem_init();
    
    int my_pe = nvshmem_my_pe();
    int n_pes = nvshmem_n_pes();
    
    // 每个 PE 分配 4 个 int 的对称内存
    int *shared_data = (int*)nvshmem_malloc(4 * sizeof(int));
    
    // 初始化本地数据
    for (int i = 0; i < 4; i++) {
        shared_data[i] = my_pe * 10 + i;
    }
    
    // 等待所有 PE 完成初始化
    nvshmem_barrier_all();
    
    // PE 0 读取 PE 1 的数据
    if (my_pe == 0) {
        int remote_val;
        nvshmem_int_get(&remote_val, &shared_data[0], 1, 1);
        printf("PE 0 read value %d from PE 1\n", remote_val);
    }
    
    nvshmem_free(shared_data);
    nvshmem_finalize();
    return 0;
}

这个例子展示了：每个 PE 初始化自己的对称内存，然后通过 nvshmem_int_get 远程读取其他 PE 的数据。注意那个 nvshmem_barrier_all——它确保所有 PE 都完成了初始化，才进行后续操作。这个同步机制我们后面会详细讲。

对称内存的底层原理

为了让你更直观地理解对称内存的布局，我画了一张图：

从这张图可以看得很清楚：每个 PE 都有自己独立的一块对称内存，但它们的虚拟地址是完全一样的。这就是 NVSHMEM 能够实现“透明远程访问”的关键——你只需要知道对方的 PE 编号和本地地址，就能直接读写对方的内存。

避坑指南

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

对称内存不能太大：每个 GPU 的显存是有限的。我曾经试图分配 16GB 的对称内存，结果在 8GB 显存的 GPU 上直接 OOM 了。
别忘了同步：远程写入后，如果不做同步，对方 PE 可能读不到最新数据。我吃过这个亏，调试了一整天才发现是同步问题。
地址对齐很重要：NVSHMEM 对对称内存的对齐有要求，尤其是做原子操作时。我建议用 nvshmem_malloc 而不是自己手动对齐。

好了，对称内存和 PE 的概念就讲到这里。记住一句话：PE 是“谁”，对称内存是“在哪”。搞清楚了这两个，后面的内容就顺了。

📚 课程目录

对称内存与 PE 概念：NVSHMEM 的基石

什么是 PE？说白了就是“干活的人”

对称内存：大家都能看到的“共享黑板”

my_pe 与 n_pes 的使用

对称内存的分配与释放

对称内存的访问模式

一个完整的例子

对称内存的底层原理

避坑指南