对称内存管理：对称内存分配、初始化、数据放置策略

对称内存，说白了就是NVSHMEM的命根子。你想想看，如果没有对称内存，PE之间怎么高效地通信？数据放哪都不对，每次访问还得绕一大圈。我刚开始接触NVSHMEM时，就在这上面栽过跟头——分配内存时没注意对称性，结果数据同步死活不对，查了两天才发现是内存地址没对齐。

嗯，今天我们就来彻底搞懂对称内存。从分配、初始化，到数据怎么放最合理，一条龙讲清楚。

4.1 对称内存的核心概念

什么叫对称内存？说白了就是：每个PE上分配一块同样大小、同样起始地址的显存。注意，是每个PE都要有，而且地址要一样。

为什么地址必须一样？因为NVSHMEM的通信机制依赖这个。比如PE 0想读取PE 1上的某个变量，它直接拿自己的对称地址去访问就行，底层会自动映射到PE 1的对应位置。如果地址不对称，这个映射就乱套了。

我在项目中遇到过一种情况：用cudaMalloc分配的内存，然后试图传给NVSHMEM用。结果呢？PE 0和PE 1的地址根本不一样，NVSHMEM直接报错。所以记住，对称内存必须用NVSHMEM提供的API来分配。

4.2 对称内存分配API

NVSHMEM提供了两套分配接口，一套是标准版，一套是可扩展版。我个人习惯用可扩展版，因为在大规模集群上更灵活。

4.2.1 nvshmem_malloc

这是最基础的分配方式。用法跟cudaMalloc差不多，但注意它分配的是对称内存。

// 每个PE分配256MB对称内存
void *ptr = nvshmem_malloc(256 * 1024 * 1024);
if (ptr == NULL) {
    printf("PE %d: 分配失败\n", nvshmem_my_pe());
}

这里有个坑：nvshmem_malloc返回的地址，在所有PE上都是一样的。你可以打印出来验证一下：

printf("PE %d: 对称地址 = %p\n", nvshmem_my_pe(), ptr);
// 你会发现所有PE打印的地址相同

4.2.2 nvshmem_calloc

如果你想要分配并清零，用这个：

// 分配并初始化为0
void *ptr = nvshmem_calloc(256 * 1024 * 1024);
// 等价于 nvshmem_malloc + memset(0)

我个人建议，如果数据需要初始化为0，直接用nvshmem_calloc。它比先malloc再memset要高效，因为底层可能用了更优化的清零方式。

4.2.3 nvshmem_realloc

这个用得少，但偶尔有用。比如你一开始分配了128MB，跑着跑着发现不够用，可以动态扩容：

void *new_ptr = nvshmem_realloc(ptr, 256 * 1024 * 1024);
if (new_ptr != NULL) {
    ptr = new_ptr;
}

注意：realloc之后，所有PE上的地址都会变。如果你之前把地址传给了别的PE，记得要同步更新。

4.3 对称内存的初始化策略

分配完内存，接下来就是初始化。这里有几个策略，我按推荐程度排序：

策略	方法	适用场景	性能
1. 本地初始化	每个PE初始化自己的那部分	数据分布均匀	高
2. 主PE广播	PE 0初始化后广播给所有人	数据相同	中
3. 分布式初始化	每个PE只初始化自己负责的块	数据有分区	高
4. 懒初始化	用到哪块就初始化哪块	稀疏数据	视情况

策略1：本地初始化——最简单，每个PE自己写自己的对称内存区域：

int *data = (int*)nvshmem_malloc(N * sizeof(int));
for (int i = 0; i < N; i++) {
    data[i] = nvshmem_my_pe();  // 每个PE填自己的编号
}

策略2：主PE广播——适合所有PE需要相同初始数据的场景：

int *data = (int*)nvshmem_malloc(N * sizeof(int));
if (nvshmem_my_pe() == 0) {
    for (int i = 0; i < N; i++) data[i] = i;
}
nvshmem_barrier_all();  // 确保PE 0写完了
// 然后用nvshmem_put或广播把数据同步给其他人

4.4 数据放置策略

数据怎么放，直接影响通信效率。这里我总结了几条黄金法则：

4.4.1 就近放置原则

数据尽量放在使用它的PE上。比如PE 0要频繁读写某个数组，那就让PE 0来分配和初始化这个数组。别让PE 3分配了，然后PE 0每次都要远程访问。

// 不好的做法：PE 3分配，PE 0频繁访问
if (nvshmem_my_pe() == 3) {
    data = nvshmem_malloc(...);
}

// 好的做法：谁用得多，谁分配
if (nvshmem_my_pe() == 0) {
    data = nvshmem_malloc(...);
}

4.4.2 数据分块策略

对于大规模数组，通常按PE数量分块。每个PE负责一块，这样通信只发生在边界区域。

int pe_id = nvshmem_my_pe();
int npes = nvshmem_n_pes();
int block_size = N / npes;
int start = pe_id * block_size;

// 每个PE只初始化自己的块
for (int i = start; i < start + block_size; i++) {
    data[i] = compute(i);
}

4.4.3 避免频繁的远程访问

如果一个PE要频繁读取另一个PE的数据，考虑把数据复制到本地。虽然复制有开销，但比起每次远程访问的延迟，复制往往更划算。

4.5 对称内存的释放

有分配就有释放。用nvshmem_free：

nvshmem_free(ptr);
ptr = NULL;  // 好习惯，防止野指针

注意：所有PE都必须调用nvshmem_free，而且传入的地址必须相同。如果一个PE忘了释放，内存泄漏是小事，更严重的是可能导致后续分配失败。

4.6 对称内存的调试技巧

调试对称内存问题，我一般用这几招：

1. 打印地址：每个PE打印自己的对称地址，确认所有PE地址一致。
2. 写后即读：分配后写一个固定值，然后所有PE读取验证。
3. 边界检查：确保没有越界访问，否则会破坏其他PE的数据。
4. 使用nvshmem_ptr：这个函数可以把对称地址转换成本地指针，方便调试。

// 调试示例：验证对称性
int *sym_addr = (int*)nvshmem_malloc(1024);
int *local_addr = (int*)nvshmem_ptr(sym_addr, nvshmem_my_pe());
printf("PE %d: sym=%p, local=%p\n", nvshmem_my_pe(), sym_addr, local_addr);
// 正常情况下，sym_addr和local_addr应该相同

4.7 知识体系总览

下面这张图总结了对称内存管理的核心脉络：

嗯，对称内存管理就讲到这里。记住三个关键词：对称分配、合理初始化、就近放置。把这三点吃透了，NVSHMEM的内存管理你就掌握了八成。

---

公众号:蓝海资料掘金营，微信deep3321