对称内存分配 API 详解：nvshmem_malloc、nvshmem_calloc、nvshmem_realloc、nvshmem_free
好，咱们直接进入正题。对称内存的分配与释放，是 NVSHMEM 编程中最基础、也最容易踩坑的地方。说白了，你写的每个 PE（处理单元）都要能访问同一块虚拟地址空间，这跟咱们平时用 malloc 完全不是一回事。
我个人习惯把对称内存分配 API 分成两类：基础分配和高级分配。基础的就是 nvshmem_malloc 和 nvshmem_free，高级的则是 nvshmem_calloc 和 nvshmem_realloc。咱们一个一个来拆。
1. nvshmem_malloc：最常用的对称分配器
这个函数是所有 NVSHMEM 程序的起点。它的签名很简单：
void *nvshmem_malloc(size_t size);你传入一个字节数，它返回一个指向对称堆的指针。注意，这个指针在所有 PE 上都是相同的虚拟地址。嗯，这里要注意：它返回的是设备指针，不是主机指针。你不能直接在 CPU 上解引用它。
我在项目中遇到过一个问题：有同事把 nvshmem_malloc 返回的指针直接传给 cudaMemcpy，结果段错误。为什么？因为 nvshmem_malloc 分配的是 NVSHMEM 管理的堆，不是普通的 CUDA 设备内存。你得用 nvshmem_ptr 或者 nvshmem_put/nvshmem_get 来访问。
核心要点：对称内存的地址在所有 PE 上一致，但数据内容可以不同。每个 PE 看到的是自己的本地副本。
2. nvshmem_calloc：分配并清零
这个函数我特别喜欢用，尤其是在初始化大型数组时。它的原型：
void *nvshmem_calloc(size_t nelem, size_t elsize);它分配 nelem * elsize 字节的对称内存，并且自动清零。你想想看，如果手动用 nvshmem_malloc 再写个循环清零，多麻烦？而且容易出错。
我曾经在调试一个分布式矩阵乘法时，发现结果总是不对。排查了半天，原来是某个 PE 上的数组没有初始化，残留了上一轮计算的数据。换成 nvshmem_calloc 后，问题瞬间解决。所以我的建议是：能用 calloc 就别用 malloc，除非你确定马上会覆盖所有数据。
小技巧：如果你分配的是大型数组（比如超过 1GB），nvshmem_calloc 的初始化开销可能很明显。这时候可以考虑用 nvshmem_malloc 配合异步清零，比如用 CUDA 内核并行清零。
3. nvshmem_realloc：重新调整大小
这个函数用得相对少一些，但在动态数据结构中很关键。它的签名：
void *nvshmem_realloc(void *ptr, size_t size);它调整之前分配的对称内存块的大小。如果新 size 大于旧 size，新增的部分不保证清零。嗯，这里要特别小心：nvshmem_realloc 可能会移动数据，返回一个新的指针。所有 PE 上的指针都会同步更新。
我记得有一次在实现动态负载均衡时，需要频繁调整每个 PE 上的工作队列大小。一开始我用 nvshmem_malloc + nvshmem_free 手动管理，代码又臭又长。后来换成 nvshmem_realloc，清爽多了。但要注意：realloc 的开销不小，因为它需要全局同步。
避坑指南：千万不要在 nvshmem_realloc 之后继续使用旧的指针。我曾经犯过这个错，结果数据全乱了。记住：realloc 后，旧指针就失效了，必须用新返回的指针。
4. nvshmem_free：释放对称内存
有分配就有释放。nvshmem_free 的用法跟 free 一样：
void nvshmem_free(void *ptr);它释放之前由 nvshmem_malloc、nvshmem_calloc 或 nvshmem_realloc 分配的内存。注意：所有 PE 都必须调用 nvshmem_free，而且传入的指针必须是对称的。如果某个 PE 忘了调用，就会造成内存泄漏。
我个人习惯在程序结束时，用一个专门的函数来统一释放所有对称内存。这样不容易漏掉。另外，nvshmem_free 是阻塞操作，它会等待所有 PE 都完成释放后才返回。
5. 各 API 对比一览
| 函数 | 功能 | 初始化 | 是否可调整大小 | 使用频率 |
|---|---|---|---|---|
| nvshmem_malloc | 分配对称内存 | 不初始化 | 否 | 最高 |
| nvshmem_calloc | 分配并清零 | 清零 | 否 | 高 |
| nvshmem_realloc | 调整大小 | 新增部分不初始化 | 是 | 中 |
| nvshmem_free | 释放对称内存 | — | — | 最高 |
6. 对称内存分配的核心逻辑
为了让你更直观地理解这些 API 之间的关系，我画了一张图。它展示了从分配、使用到释放的完整流程，以及每个 API 在其中的位置。
从图中你可以看到，无论你选择哪个分配 API，最终都要经过 nvshmem_free 来释放。而且，nvshmem_realloc 其实是一个特例——它既可以分配新内存，也可以释放旧内存，但它的使用场景比较受限。
7. 避坑总结
- 不要混用分配器：用
nvshmem_malloc分配的内存，必须用nvshmem_free释放。不能用cudaFree或free。 - 所有 PE 必须同步：分配和释放操作在所有 PE 上都要执行，否则会导致死锁或内存泄漏。
- 检查返回值：如果分配失败，返回 NULL。我见过不少代码不检查返回值，结果在后续操作中崩溃。
- 对称内存不是持久化的：程序结束后，对称内存自动释放。不要指望它跨程序运行。
好了，对称内存分配 API 就讲到这里。这些函数虽然简单，但用好了能避免很多坑。下一章咱们聊聊对称内存的访问模式——怎么高效地读写这些数据。