3. Team 创建 API：nvshmem_team_create_strided 与 nvshmem_team_create_2d 详解
好，咱们接着聊 Team 的创建。上一节讲了最基础的 nvshmem_team_split_strided，说白了就是按固定步长切分。但实际项目中，你遇到的拓扑结构往往更复杂。比如，你要做 2D 网格通信，或者需要把 GPU 按某种非均匀的 stride 分组。
这时候，nvshmem_team_create_strided 和 nvshmem_team_create_2d 就派上用场了。我个人习惯把前者叫做「自由步长版」，后者叫做「网格拓扑版」。它们俩的底层逻辑完全不同，但都是构建复杂通信模式的核心工具。
3.1 nvshmem_team_create_strided：自由步长分组
这个 API 的签名长这样：
int nvshmem_team_create_strided(
nvshmem_team_t parent_team,
int start,
int stride,
int size,
nvshmem_team_config_t *config,
nvshmem_team_t *new_team
);参数不多，但每个都有讲究：
- parent_team：父团队，通常是
NVSHMEM_TEAM_WORLD或者你之前 split 出来的子团队。 - start：起始 PE 的全局索引。注意，这个索引是相对于父团队的。
- stride：步长。相邻两个 PE 在父团队中的索引差。
- size：新团队包含的 PE 数量。
- config：配置结构体，一般传
NULL用默认值。 - new_team：输出参数，新团队的句柄。
举个例子，假设你有 8 个 PE（0~7），你想挑出索引为 1、3、5、7 的四个 PE 组成一个新团队：
nvshmem_team_t strided_team;
int start = 1;
int stride = 2;
int size = 4;
nvshmem_team_create_strided(
NVSHMEM_TEAM_WORLD,
start,
stride,
size,
NULL,
&strided_team
);嗯，这里要注意：新团队内部的 PE 排名（rank）是按照你选取的顺序来的。start=1 的 PE 在新团队里是 rank 0，start+stride=3 的是 rank 1，以此类推。
3.2 nvshmem_team_create_2d：网格拓扑创建
这个 API 就更有意思了。它专门用来创建 2D 网格状的团队结构。你想想看，很多科学计算、图像处理、矩阵运算，天然就是 2D 网格的通信模式。
函数签名：
int nvshmem_team_create_2d(
nvshmem_team_t parent_team,
int xrange,
int yrange,
nvshmem_team_config_t *config,
nvshmem_team_t *new_team_x,
nvshmem_team_t *new_team_y
);参数解析：
- parent_team：父团队。
- xrange：X 方向的维度大小。
- yrange：Y 方向的维度大小。注意，xrange * yrange 必须等于父团队的 PE 总数。
- config：配置，传 NULL 即可。
- new_team_x：输出参数，X 方向的行团队。每行包含 yrange 个 PE。
- new_team_y：输出参数，Y 方向的列团队。每列包含 xrange 个 PE。
它的核心逻辑是：把父团队的 PE 按行优先（row-major）排列成一个 xrange 行、yrange 列的网格。然后，每一行形成一个行团队（X 方向），每一列形成一个列团队（Y 方向）。
举个例子，12 个 PE（0~11），你想排成 3 行 4 列：
nvshmem_team_t row_team, col_team;
int xrange = 3; // 行数
int yrange = 4; // 列数
nvshmem_team_create_2d(
NVSHMEM_TEAM_WORLD,
xrange,
yrange,
NULL,
&row_team,
&col_team
);创建之后，PE 0、1、2、3 属于同一个行团队（第 0 行），PE 0、4、8 属于同一个列团队（第 0 列）。
3.3 两个 API 的对比与选择
| 特性 | nvshmem_team_create_strided |
nvshmem_team_create_2d |
|---|---|---|
| 分组方式 | 按固定步长选取 PE | 按 2D 网格行列分组 |
| 输出团队数 | 1 个 | 2 个（行团队 + 列团队） |
| 适用场景 | 非均匀采样、环形拓扑 | 矩阵运算、图像处理、物理模拟 |
| 灵活性 | 高，步长可任意指定 | 低，必须满足网格约束 |
| 内部 rank 映射 | 按选取顺序映射 | 行团队按行内顺序，列团队按列内顺序 |
说白了，如果你的通信模式是「每隔几个 PE 取一个」，用 strided。如果你的通信模式是「按行或按列做集体操作」，用 2d。我在项目中遇到过好几次，有人用 strided 硬生生模拟 2D 网格，结果代码又长又容易出错。其实直接用 nvshmem_team_create_2d 就完事了。
3.4 实战：结合 2D 团队做矩阵转置
咱们来看一个实际例子。假设你要做矩阵转置，数据分布在 2D 网格上。每个 PE 持有矩阵的一个子块。转置操作需要每个 PE 跟同行同列的 PE 交换数据。
// 假设 4 个 PE，排成 2x2 网格
nvshmem_team_t row_team, col_team;
nvshmem_team_create_2d(NVSHMEM_TEAM_WORLD, 2, 2, NULL, &row_team, &col_team);
// 获取行团队和列团队中的 rank
int row_rank = nvshmem_team_my_pe(row_team);
int col_rank = nvshmem_team_my_pe(col_team);
// 行团队内做 alltoall：交换行内数据
nvshmem_team_alltoall(sendbuf, recvbuf, sizeof(float) * block_size, row_team);
// 列团队内做 alltoall：交换列内数据
nvshmem_team_alltoall(sendbuf, recvbuf, sizeof(float) * block_size, col_team);你看，两行代码就完成了转置的核心通信。如果没有 2D 团队，你得手动计算每个 PE 的通信对，代码量至少翻三倍。
nvshmem_team_my_pe 获取当前 PE 在行/列团队中的局部 rank。这个 rank 在后续的集体操作中非常有用，比如决定当前 PE 负责哪一块数据。
3.5 避坑指南
- 团队数量限制： NVSHMEM 对同时存在的团队数量有限制（通常是 8 个左右）。不要创建了不销毁，否则会耗尽资源。我见过有人循环创建团队，结果跑到第 9 次就报错了。
- 销毁顺序： 先销毁子团队，再销毁父团队。否则父团队销毁后，子团队句柄就变成野指针了。
- 配置参数：
nvshmem_team_config_t里有个split_priority字段，可以用来控制团队分裂时的优先级。默认值就够了，除非你有特殊的负载均衡需求。 - 性能考量： 创建团队本身是同步操作，所有 PE 必须同时调用。如果某个 PE 调用晚了，整个程序会卡住。所以，团队创建最好在初始化阶段完成，不要在热路径上频繁创建销毁。
好了，关于这两个 API 的用法就讲到这里。记住，strided 适合自由步长，2d 适合网格拓扑。选对了工具，代码写起来顺手，性能也更容易优化。