AllGather ND 环

此页面上的所有地址适用于 libtpu-0.0.40-cp314 轮中的 libtpu.so(内部版本号 89edbbe81c5b328a958fe628a9f2207d，内部版本 libtpu_lts_20260413_b_RC00)。其他版本会有所不同。 .text VMA等于文件偏移量；所有地址均为 VMA。

摘要

ND 环 AllGather 是水母后端如何将 HLO all-gather(操作码 6)通过 1-D/2-D/3-D 物理环面转换为 ICI 环流量的方式。与全归约系列(在多个活动轴上分解为归约分散 + 全收集阶段)不同，AllGather 发射器是明确的轴分解：它为每个网格轴具体化一个每轴 RingLocation 表，并且在发射时逐轴遍历环，在每个 DMA 阶段推进一个网格环。执行此操作的读取器阶段是一个干净的两阶段组合：GetShardIndex 在预先计算的 RingLocation 状态上使用单个模块化步骤 (base ± step) mod ring_len 计算每个轴的每步环坐标，然后 GetOffset 通过从小到大混合基数点积将完整的每轴坐标向量线性化为平坦的聚集缓冲区槽。这是 MPI 环全聚集的 TPU 模拟，推广为矩形环面，环旋转表示为设备序号上的模算术，而不是显式的邻居遍历。

选择要响铃的“多少”轴不是单个立即数。它是一个合词：TpuCompilationEnvironment 使能字节(+0xbe 表示 2 维，+0xc0 表示 3 维)和集体的设备列表是否以几何方式投影到 2 轴或 3 轴平面，由 ReplicaGroupsOnNDPlane(plane_dim=2/3) 测试。 ReplicaGroupsOnNDPlane 是每轴 MeshNDInfo 矢量的构建器；如果设备列表适合 plane_dim 轴平面，则投影成功。然后，在安装读取器消耗的环序和每维大小向量之前，Init2DAllGather / Init3DAllGather 通过计算 MeshNDInfo+0x38 中的活动轴位(popcount 2 或 3)来重新验证投影。

此页面拥有 ND 环 AllGather 分片数学 (GetShardIndex / GetOffset / ComputeAdjustedIndexAtRuntime)、2-D / 3-D 网格维度选择器 (UseAllGather2D / UseAllGather3D)、每轴环安装(Init{1,2,3}DAllGather→InitDim)，以及**MeshNDInfo环几何**读者索引。每轴设备的一代→邻居表 (CreateStaticNDRingReplicaInfoTable) 位于 常量映射器 上；在这里进行全聚集的策略选择位于SelectND策略；戒指的成本位于SPMD 链路计数成本。此页面链接这些内容并且不会重新派生它们。

重新实现，合约为：

• 分片数学。 GetShardIndex 读取 RingLocation[dim] 和 ring_len = sizes[dim]，将每个轴的环基提升到坐标向量中，将该轴的步进坐标计算为 (base ± step) mod ring_len(由布尔选择的 bidir 与向前)，将非行走轴固定到序数 0，并且调用 GetOffset 进行重新线性化。
• 线性化器。 GetOffset 是从小到大混合基数点积 offset = (Σ_k coords[mtm[k]] · Π_{j<k} bounds[mtm[j]]) mod Π bounds，由三跨度等长 RET_CHECK 保护。
• 选择器。 1-D / 2-D / 3-D = (env enable byte) ∧ (ReplicaGroupsOnNDPlane projects onto a 2/3-axis plane)，其中 MeshNDInfo+0x38 popcount = #active mesh axes 在安装时重新验证。

分片索引读取器	xla::jellyfish::(anonymous namespace)::GetShardIndex @ 0x13811600
缓冲线性化器	xla::jellyfish::(anonymous namespace)::GetOffset @ 0x138106c0 (all_gather_emitter.cc:164)
短环重新调整	ComputeAdjustedIndexAtRuntime @ 0x13800d00
2-D 选择器	AllGatherEmitter::UseAllGather2D @ 0x13801740(环境 +0xbe、plane_dim=2)
3-D 选择器	AllGatherEmitter::UseAllGather3D @ 0x13801a40(环境 +0xc0、plane_dim=3)
平面投影仪	ReplicaGroupsOnNDPlane @ 0x1c890960 → optional<vector<MeshNDInfo>>
每轴安装	Init1DAllGather @ 0x13807180 / Init2DAllGather @ 0x13807720 / Init3DAllGather @ 0x13807aa0
每轴环填充	AllGatherEmitter::InitDim @ 0x13804980 → RingLocation[] 在 this+0x408
MeshNDInfo 尺寸	0x40 B(复制向量 0x127b5100) — 轴 id 向量 / 大小向量 / 环序向量 / 暗淡位掩码
RingLocation 步幅	0x38 B(7 个 qword)
信心	HIGH(GetShardIndex、GetOffset、UseAllGather2D、Init2DAllGather 的反编译验证主体)，除非行/标注另有说明

---

这在哪里

SelectNDStrategy (SelectND策略) 决定“全聚集”实现为 ND 环；此页面是所选环实际运行的发送时读取管道。构建侧 — CreateStaticNDRingReplicaInfoTable 将每个网格轴注册一个 device_id → ring-neighbor 表作为 ConstantMapper 类型 0/1/2 — 位于 常量映射器 上。此页面是消费者：决定要安装多少个表的维度选择器，从表中具体化每个轴的 RingLocation 的 InitDim 行走，以及将环位置转换为平面聚集缓冲区槽的 GetShardIndex / GetOffset 算术。

管道，端到端：

UseAllGather2D / UseAllGather3D          (dimensionality select: env byte ∧ plane projection)
        │   ReplicaGroupsOnNDPlane(plane_dim=2/3) → optional<vector<MeshNDInfo>>
        ▼
Init{1,2,3}DAllGather                     (popcount-verify MeshNDInfo+0x38, install ring-order + sizes)
        │   InitDim once per mesh axis  → RingLocation into this+0x408[dim]
        ▼
Phase{Zero,One,Two}DmaNDPlaneAllDimensionsStart   (one mesh ring advanced per DMA phase)
        │
        ▼
GetShardIndex (per axis)  ── (base ± step) mod ring_len  over RingLocation[dim]
        │
        ▼
GetOffset                 ── minor-to-major mixed-radix linearization → flat buffer slot
        │   (ComputeAdjustedIndexAtRuntime rescales when this ring axis is shorter than the longest)
        ▼
ICI DMA descriptor (PerBufferDmaEmitter)
```text

---

## MeshNDInfo — ND 环几何形状

环几何结构在 `MeshNDInfo`(0x40 字节；复制向量 `0x127b5100`)中承载。一个 `MeshNDInfo` 描述一个*平面*：其活动网格轴、沿每个环的环长度、沿每个环的设备序号以及其中环面轴参与的位掩码。 `ReplicaGroupsOnNDPlane` 构建 `vector<MeshNDInfo>` — `plane_dim` 条目 — 并且 AllGather 安装程序读取其元素 `[0]`。

| 字段 | 偏移 | 型 | 含义 |
|---|---|---|---|
| `minor_to_major` | `+0x00` | `vector<MeshDim>` (`int32×N`) | 网格轴 ID 列表，按小到大(基数)顺序； `size` 在 `+0x08` |
| `dim_sizes` | `+0x18` | `vector<long>` | 每个维度**环长度**(每轴一个)； `size` 在 `+0x20` |
| `ring_order` | `+0x28` | `vector<MeshDim>` (`int32×N`) | **环遍历顺序** — 沿每个环的设备序号； `size` 在 `+0x30` |
| `dim_bitmask` | `+0x38` | `long` | **活动环面轴**的低 3 位位掩码； `popcount(bitmask & 7)` = `#axes` |

`+0x38` 处的 `dim_bitmask` 是维度预言机。 `Is2D()` 为 `popcount(bitmask & 7) == 2`； `Is3D()` 是 `popcount(bitmask & 7) == 3`。 `CreateStaticNDRingReplicaInfoTable` ([常量映射器](constant-mapper.md)) `RET_CHECK`s `mesh_info.Is2D() || mesh_info.Is3D()` (`net_util.cc:2440`)，安装人员重新验证它(下一节)。

> **注意 —** `MeshNDInfo` 三元组 — 轴 ID 向量、环长度向量、环阶向量 — 一对一映射到三个参数范围 `GetShardIndex` / `GetOffset` 消耗：`minor_to_major` 成为线性化器的基阶， `dim_sizes` 变为 `bounds`(模块化环长度)，而 `ring_order` 是嵌入到每个 `RingLocation` 中的每轴序号查找。重新实现者应该将 `MeshNDInfo` 视为“一个 ND 环的基分解”。
>
> **GOTCHA —** `+0x38`内部的位→轴映射(该位是网格轴0/1/2)是与`+0x00` axis-id向量一致的结构读取；只有 popcount→dimension-count 语义是字节确认的(`Is2D`/`Is3D` 检查中的 popcnt 指令)。将维数从弹出计数中驱动出来，而不是从假定的位位置中驱动。 (每位分配为中等，弹出计数为高。)

---

## 2-D / 3-D 选择器

### 用途

`UseAllGather2D` 和 `UseAllGather3D` 回答“2 轴还是 3 轴上的环？”每个都是环境启用字节和几何平面投影测试的结合。它们在 `AllGatherEmitter::GenerateConstants` (`0x13801be0`) 中进行尝试：首先是 3-D (`0x13801dfd`)，然后是 2-D (`0x1380203f`)，否则是 1-D 回退；显式路由路径 (`ShouldUseExplicitRouting` @ `0x13803aa0`) 相反会构建 Type-5 路由表(请参阅 [路由](../routing/overview.md))。

### 算法

```c
// AllGatherEmitter::UseAllGather2D — 0x13801740. Returns {projected_groups, has_value} (0x18 B).
function UseAllGather2D(target, hlo, dev_assign, env, bidir):
    // ENABLE GATE
    if env == nullptr:
        if hlo->GetModule() == nullptr: return {_, false}
        env = GetTpuCompEnv(hlo, target)
    if IsPartOfCollectiveComputeFusion(hlo, target): return {_, false}   // 0x1380176f
    if env[0xbe] == 0:                  return {_, false}   // 2-D enable byte (190)
    if dev_assign == nullptr:                                // fall back to static DA
        if module_config[+0x660] == 0:  return {_, false}    // static_device_assignment.has_value()
        dev_assign = HloModuleConfig::static_device_assignment(...)
    if target->vtable[0x18](target):    return {_, false}    // a Target capability veto

    // PROJECTION GATE
    device_list = GetCollectiveDeviceList(hlo, dev_assign)              // 0x13801871
    planes = ReplicaGroupsOnNDPlane(target, dev_assign, device_list,
                                    /*plane_dim=*/2, bidir)             // 0x138018a3 (mov $0x2,%r8d)
    if !planes.ok:                      return {_, false}               // success byte at result+0x18
    X = planes[0].dim_sizes[0]; Y = planes[0].dim_sizes[1]             // result+0x18 reads
    if env[0xbf] == 0 && hlo->opcode_field == 8                         // continuation-fusion arm
       && !ShouldUseContinuationFusionAllGather(...) && X != Y:
        return {_, false}                                              // square-plane requirement
    return {planes, true}                                             // has_value byte at result+0x10

UseAllGather3D (0x13801a40) 的字节形状相同，有两次替换：使能字节为 env[0xc0] (0x13801a65)，投影为 ReplicaGroupsOnNDPlane(..., plane_dim=3) (0x13801b52, mov $0x3,%r8d);它的投影成功字节位于与 2D 版本相同的帧槽(cmpb $0x1,-0x38(%rbp) @ 0x13801b5d，即 optional 的 has_value 字节)，并且其连续融合覆盖读取 env 字节 env[0xc1](cmpb $0x0,0xc1(%rax) @ 0x13801b6a) — 2D env[0xbf] 臂的 3D 模拟。

选择器	使能字节	投影	成功/形状测试	初始化路径
UseAllGather2D	TpuCompEnv[+0xbe]	ReplicaGroupsOnNDPlane(plane_dim=2)	has_value字节@rsp+0x18(cmpb $0x1@0x138018ae)；由 env +0xbf 门控的方形平面 X==Y 臂	Init2DAllGather(2轴)
UseAllGather3D	TpuCompEnv[+0xc0]	ReplicaGroupsOnNDPlane(plane_dim=3)	has_value字节@rsp+0x18(cmpb $0x1@0x13801b5d)；由 env +0xc1 门控的连续融合臂	Init3DAllGather(3轴)

QUIRK — 2-D 臂施加 方形平面 要求：当不采用连续融合路径时，它会拒绝 X != Y(两个投影环长度必须匹配)。无条件允许矩形二维平面的重新实现者将在将二进制路径塑造为一维时过度选择二维环。连续融合覆盖(ShouldUseContinuationFusionAllGather)是有记录的逃生舱口； (ShapeSize & 3) == 0 对齐探针(0x13801740 主体)还会在投影之前强制设备列表路径。

明白了 — ReplicaGroupsOnNDPlane 不是谓词 — 它是 MeshNDInfo 向量的 构建器，使用所需的 plane_dim 进行调用。它的 plane_dim 参数等于它返回的 MeshNDInfo 计数，它等于 ND 环表的数量和安装程序重新验证的 popcount。将“决定维度”从“构建环表”分成两遍的重新实现将与融合它们的二进制文件不同：投影是一次调用中的决策和构造。它存储在 NDPlaneCacheKey → optional<vector<MeshNDInfo>> 缓存(0x225799b8)上，因此使用相同的 plane_dim 重新查询相同的设备列表是免费的；缓存键是设备分配字符串 ⊕ 序列化拓扑 (TpuTopologySerdes::Distill → TpuTopologyArgs::ToProto → SerializeToStringDeterministic) ⊕ plane_dim int ⊕ bool，全部在 nd_plane_cache_mutex 下。实际的坐标投影数学在 ReplicaGroupsOnNDPlane 中不是内联 — 入口函数仅构建密钥、获取锁、调度，并且在未命中时调用每维 lambda ReplicaGroupsOnNDPlaneImpl::$_0 (0x1c896400) 每个网格轴一次 (plane_dim ∈ {0,1,2}，传递为0x100000000 | axis)。那个 lambda 就是投影所在的地方；它的主体在 张量分割/ND 平面 页面的密集 ReplicaGroupsOnNDPlane 投影部分进行反编译。

---

每轴环安装

用途

一旦选择了 MeshNDInfo 矢量，Init{1,2,3}DAllGather 安装环：重新验证维数，每个网格轴驱动 InitDim 一次以具体化每个轴的 RingLocation，并复制轴顺序矢量(m0.minor_to_major → this+0x218)和每个尺寸的向量(m0.dim_sizes → this+0x1e8)进入发射器，以便读者可以索引它们。安装程序不读取MeshNDInfo自己的ring_order字段(m0+0x28)； this+0x218 处的发射器阶向量是根据 minor_to_major 构建的。

算法

// AllGatherEmitter::Init2DAllGather — 0x13807720. (Init3D @0x13807aa0 differs only in counts/popcount.)
function Init2DAllGather(meshes /*Span<MeshNDInfo>*/, hlo):
    RET_CHECK(!meshes.empty())                       // "mesh_info should have size larger than 1" (line 2970)
    if InitHierarchicalStates(...) != ok: return err  // line 2972
    this->mesh_axes /*this+0x1d0*/ = {axis0, axis1}    // assign 2 MeshDim ids
    for k in [0, this->mesh_axes.size):               // 0x1380779d, stride 4
        if InitDim(hlo, mesh_axes[k], /*reorder=*/0) != ok: return err   // line 2976
    this->initialized_2d /*this+0x19c*/ = 1            // 0x138077b5

    m0 = meshes[0]
    RET_CHECK(m0.minor_to_major[0] == y || == x)      // "could only be x or y" (line 2980)
    RET_CHECK(m0.minor_to_major.size == 2              // m0+0x08
              && m0.dim_sizes.size == 2                // m0+0x20  (cmpq $0x2,0x20(%r14) @0x138077e3)
              && popcount(m0.dim_bitmask & 7) == 2)    // m0+0x38 — "mesh_info[0].Is2D()" (line 2982)
    this->ring_order /*this+0x218*/ = m0.minor_to_major // from m0+0x00, count 2 (0x13807819)
    this->dim_sizes  /*this+0x1e8*/ = m0.dim_sizes     // from m0+0x18, count 2 (0x13807838)
    return ok
```text

`Init3DAllGather` (`0x13807aa0`) 与 count-3 替换相同：`this+0x19d=1`(3D 初始化)，安装复制 count-3 `minor_to_major`(到 `this+0x218`)和 `dim_sizes`(到`this+0x1e8`)，popcount 检查实现为 `not(bitmask); test $7`(所有低 3 位设置⇔ popcount 3)，而不是显式的 `popcnt == 3`(`mov 0x38(%r14),%eax; not %eax; test $0x7,%al` @ `0x13807b6e`)。 `Init1DAllGather`(`0x13807180`)调用`InitHierarchicalStates`，分配单元素轴列表，并与`reorder=1`一起运行一次`InitDim`。

> **注意 —** `Init2D`/`Init3D` 内部的冗余 popcount 重新验证是故意的，而不是偏执：`ReplicaGroupsOnNDPlane` 的结果被缓存并且可以在选择器调用之间重用，因此安装程序断言缓存的 `MeshNDInfo` 实际上与它即将要的维度匹配安装。两个 `RET_CHECK` 是不同的事实 — `minor_to_major[0] ∈ {x, y}`(短轴是真正的环面轴，第 2980 行)和 `Is2D()`(恰好两个活动轴，第 2982 行)。

### InitDim — RingLocation 填充

`InitDim` (`0x13804980`) 填充一根轴。它读取此 `MeshDim` 的每轴常量 — `GetConstant(MeshDim==dim)`(ND 环邻居表，类型 0/1/2；`0x13804c58`)，以及 `GetConstant(Type 3)`(静态 ND 环；`0x13804c4e`)和 `GetConstant(Type 4)`(有限 ICI 路由； `0x13804ae6`) 作为备用 — 然后在重新排序路径上调用 `net_util::GetRingLocation` (`0x1c6a0c40`) 或 `GetRingLocationWithReordering` (`0x1c6a19c0`)，并将生成的 `0x38` 字节 `RingLocation` 存储到 `this+0x408[dim]` 中。

```text
this+0x408 / +0x410 / +0x418  : vector<RingLocation> {data, size, cap}   (per-axis ring state)
  RingLocation stride 0x38 (7 qwords); RingLocation[dim].first is this axis's ring base value

this+0x408 正是 rings 跨度 GetShardIndex 索引，每个网格轴一个元素，在 Init* 驱动程序下由 InitDim 一次填充一个轴。

---

碎片数学

GetShardIndex — 每轴旋转一圈

GetShardIndex (0x13811600) 对于一环步长的一个网格轴计算本地核心应读取的平面聚集缓冲区槽。旋转是预先计算的 RingLocation 状态上的单个模块化步骤 - 在发射时没有邻居行走循环；相邻关系被烘焙到每轴表中，并且步骤是算术的。

算法

// (anon)::GetShardIndex — 0x13811600
// rings = this+0x408 (Span<RingLocation>, stride 0x38 = 7 qwords)
// sizes = MeshNDInfo dim_sizes (the ring lengths)
function GetShardIndex(rb, base /*LloValue*/, dim, bidir,
                       rings, sizes, minor_to_major, coords_span):
    ring_base = rings[dim].first        // a5[7*dim]            — this axis's ring base value
    ring_len  = sizes[dim]              // a7[2*dim]            — the modulus for this axis

    // (1) lift every axis's ring base into a coordinate vector  (grow loop @0x13811670)
    coords = vector<LloValue*>()
    for k in [0, rings.size):
        coords.push_back(rings[k].first)

    // (2) THE RING ROTATION — this axis's per-step coordinate, modulo the ring length
    if bidir:                                              // a4 != 0
        // (ring_base - base + ring_len) mod ring_len   — backward+forward step
        t = rb.SsubS32(ring_base, base)                   // 0x1381175f
        t = rb.SaddS32(t, rb.SimmS32(ring_len))           // 0x1381178e (+ring_len keeps it ≥0)
    else:
        // (ring_base + base) mod ring_len               — forward-only step
        t = rb.SaddS32(ring_base, base)                   // 0x1381178e
    stepped = rb.SmodU32(t, ring_len)                     // 0x1381179c — modulus = sizes[dim]

    // (3) write the stepped coordinate back; pin non-walked axes to ordinal 0
    coords[dim] = stepped                                 // 0x138117b9
    for slot in coords_span:                              // 0x138117d0
        coords[slot] = rb.SimmS32(0)

    // (4) relinearize the full coordinate vector into the flat buffer slot
    offset = GetOffset(minor_to_major, coords, sizes, rb) // 0x13811813
    free(coords); return offset
```text

bidir/forward split是唯一的分支：双向环计算`(ring_base − base + ring_len) mod ring_len`(`+ring_len`保证无符号模的非负参数)，单向环计算`(ring_base + base) mod ring_len`。每个其他轴都固定到序数 0，因此单个 `GetShardIndex` 调用恰好前进一圈，同时保持其余部分固定 - 这就是为什么三个 DMA 阶段各自前进不同的轴。

> **QUIRK —** 模数始终为 `sizes[dim]`(*此*轴的环长)，且 `SmodU32` 无符号。 bidir 臂中 mod 之前的 `+ring_len` 不是一种风格选择 - 它是正确性守卫：`ring_base − base` 可以为负，并且负包裹值的无符号 mod 是错误的。将 bidir 臂折叠到 `(ring_base − base) mod ring_len` 的重新实现者将为每个后退步骤生成错误的分片索引。 (高 — 在反编译中得到确认：`SsubS32`，然后是 `SaddS32(SimmS32(ring_len))`，然后是 `SmodU32`。)

### GetOffset — 小到大混合基线性化

`GetOffset` (`0x138106c0`) 将每轴坐标向量折叠到一个平槽中。它是标准的行主线性化，但采用*次到主*基数顺序(`minor_to_major` 跨度选择轴顺序和边界的运行乘积)。

### 算法

```c
// (anon)::GetOffset — 0x138106c0
function GetOffset(minor_to_major, coords, bounds, rb):
    // all three spans must be the same length
    RET_CHECK(minor_to_major.size == coords.size
              && coords.size == bounds.size)      // all_gather_emitter.cc:164
    if minor_to_major.size == 1:
        return coords[0]                          // single-axis fast path (0x138106fb)

    acc = rb.SimmS32(0)
    for k in [0, n):
        term = coords[minor_to_major[k]]
        for j in [0, k):                          // running radix product
            term = rb.SmulU32(term, rb.SimmS32(bounds[minor_to_major[j]]))
        acc = rb.SaddS32(acc, term)               // 0x1381074e
    radix = Π_k bounds[minor_to_major[k]]          // vectorized vpmulld @0x13810816
    return rb.SmodU32(acc, radix)                  // 0x13810898

即offset = (Σ_k coords[mtm[k]] · Π_{j<k} bounds[mtm[j]]) mod (Π_k bounds[mtm[k]])。最终的 mod (Π bounds) 将线性化索引包装到收集缓冲区的槽计数中。边界乘积是通过边界数组上的手动矢量化 vpmulld 缩减(4 宽，来自 dword_84A2B08 == 1)计算得出的，余数为标量尾部——这是一种不改变公式的微优化。

明白了 — RET_CHECK 证明了跨度身份 minor_to_major.size == coords.size == bounds.size；这就是告诉逆向工程这三个不透明跨度正是 {radix order, per-axis coordinates, per-axis ring lengths} 的原因。如果重新实现传递由原始轴 id 而不是 minor_to_major[j] 索引的 bounds 数组，则运行的乘积将采用错误的基数顺序，并且对于任何重要的小到大排列，偏移量都是错误的。索引 bounds 到 minor_to_major，与内循环完全相同(bounds[minor_to_major[j]]，在 0x13810790 反编译)。

ComputeAdjustedIndexAtRuntime — 短环重新缩放

GetShardIndex 的槽位是根据该轴的环长计算的；当一个网格轴比最长的轴短时，必须按比例放大索引，以便较短环的槽映射到完整的聚集缓冲区。 ComputeAdjustedIndexAtRuntime (0x13800d00) 通过谓词掩码来执行此操作：

// ComputeAdjustedIndexAtRuntime — 0x13800d00
function ComputeAdjustedIndexAtRuntime(rb, core_idx, idx, dim):
    p0 = rb.SeqS32(core_idx, 0)                       // 0x13800d3e
    p1 = rb.SeqS32(core_idx, 1)                       // 0x13800d62
    is2d = (this->mesh_axes.size /*this+0x1d8*/ == 2) // 0x13800d6c
    mask = p0 | (~is2d & p1)                           // Pimm/Pneg/Pand/Por (0x13800d7c..)
    ratio = 1                                          // default (LODWORD(v10)=1)
    if !GetTpuCompEnv[+0x13f2]:                         // 0x13800db5 — rescale enabled when byte is 0
        max_dim = max_element(this->dim_sizes)        // this+0x1e8 / size this+0x1f0 (0x13800e10..)
        ratio   = max_dim / dim_sizes[dim]            // idiv (0x13800f10)
    scaled  = rb.SdivS32(idx, ratio)                  // 0x13800f30
    return rb.Sselect(mask, idx, scaled)              // 0x13800f41 — args: (mask, idx, scaled)
```text

重新缩放将索引除以 `max_dim / ring_len[dim]`，并在每核谓词 `mask` 上发出 `Sselect`；二进制将参数作为 `Sselect(mask, idx, scaled)` 传递（在 `0x13800f3b` 验证：`%rdx`=idx，`%rcx`=scaled）。当环境字节 `+0x13f2` 被设置时，`max_element` 扫描会被跳过，`ratio` 保持 `1`，因此 scaled 分支折叠为 `idx`。它的调用者是 async window-emission 路径：`AsyncAllGatherEmitter::EmitWindow`（`0x137eec20`，调用 `0x137eed9a`）和 `AsyncAllGatherEmitter::MaybeEmitWindow`（`0x137eefa0`，调用 `0x137ef393` / `0x137ef583`），而不是 `GetPhazeZeroShardIndexHelper` / `GetShardIndex` 链。单独的 hierarchical-split 递归（`MaybeMapShardIndexForHierarchicalSplit` @ `0x138108e0`，从 `GetPhazeZeroShardIndexHelper` @ call `0x137f1818` 到达的 per-chip core-split fold）只做了部分追踪。

### 呼叫站点 — 每个 DMA 阶段响铃一次

`GetShardIndex` 有三个调用者，都在 AllGather 发射器中，每个调用者都将发射器的 `this+0x408` `RingLocation` 数组作为 `rings` 传递：

| 呼叫者 | 地址 | 角色 |
|---|---|---|
| `GetPhazeZeroShardIndexHelper` | `0x137f1780`(呼叫`0x137f1803`) | Phase-0分片索引+分层分割图 |
| `PhaseOneDmaNDPlaneAllDimensionsStart` | `0x137f42c0`(呼叫`0x137f5011`) | Phase-1环阶(单元件MeshDim) |
| `PhaseTwoDmaNDPlaneAllDimensionsStart` | `0x137f9060`(呼叫`0x137f9e6d`) | 二相环步 |

三相依次进环；确切的每轴分工(每个阶段走哪个轴，通过 `GetRecvFlags` @ `0x137f5cc0` 的发送/接收 SFLAG 相互作用，以及 `PipelinedAllGatherSlots` @ `0x137f4100` 重叠)是 DMA 循环层，这里没有追踪(它属于 [路由](../routing/overview.md) / [ICI 面料](../ici/overview.md) 部分)。

---

## 功能图

| 功能 | 地址 | 角色 |
|---|---|---|
| `(anon)::GetShardIndex` | `0x13811600` | 每轴环旋转 `(base±step) mod ring_len` |
| `(anon)::GetOffset` | `0x138106c0` | 小到大混合基线性化器 |
| `ComputeAdjustedIndexAtRuntime` | `0x13800d00` | 短环指数调整 |
| `GetPhazeZeroShardIndexHelper` | `0x137f1780` | `GetShardIndex` + 分层分割包装器 |
| `MaybeMapShardIndexForHierarchicalSplit` | `0x138108e0` | 每芯片核心分割分片重新映射 |
| `AllGatherEmitter::UseAllGather2D` | `0x13801740` | 2-D 选择器(env `+0xbe` ∧ `plane_dim=2`) |
| `AllGatherEmitter::UseAllGather3D` | `0x13801a40` | 3-D 选择器(env `+0xc0` ∧ `plane_dim=3`) |
| `ReplicaGroupsOnNDPlane` | `0x1c890960` | 平面投影仪/`MeshNDInfo`构建器(缓存) |
| `GetCollectiveDeviceList` | `0x13801940` | 用于投影的设备列表解析器 |
| `AllGatherEmitter::Init1DAllGather` | `0x13807180` | 单轴安装 |
| `AllGatherEmitter::Init2DAllGather` | `0x13807720` | 2轴安装+ popcount-2验证 |
| `AllGatherEmitter::Init3DAllGather` | `0x13807aa0` | 3轴安装+ popcount-3验证 |
| `AllGatherEmitter::InitDim` | `0x13804980` | 每轴 `RingLocation` 填写为 `this+0x408` ND 环信息表中的 |
| `net_util::GetRingLocation` | `0x1c6a0c40` | `RingLocation` |
| `MeshNDInfo` 复制器 | `0x127b5100` | 0x40-B 几何：轴 ID / 尺寸 / 环序 / 位掩码 |

---

## 未解决的问题

- **`ReplicaGroupsOnNDPlane` 内部结构。** 此页面确认它返回 `optional<vector<MeshNDInfo>>`，缓存在 `NDPlaneCacheKey` 上，并且 `plane_dim` = `MeshNDInfo` 计数 = 环维数。坐标投影本身位于每轴 lambda `ReplicaGroupsOnNDPlaneImpl::$_0` (`0x1c896400`) 中，而不是内嵌在 `0x1c890960` 条目中：每个设备的 `chip_coordinates`(通过 `TensorCoreLocationForLogicalDeviceId` `0x1c8904e0`)针对网格范围进行混合基数线性化，然后分派到`ReplicaGroupForm2DRing` (`0x1c88e6e0`) 或 `ReplicaGroupsOn3DPlane` (`0x1c8901e0`) 由 `n_dim` 组成，带有 `n_dim==1` 短路；所选轴的范围预先乘以 `LogicalDevicesPerChip`，以将子核折叠到该环中。在[张量分割/ND 平面](tensor-split-ndplane.md)上反编译。 HIGH为调度+子核心弃牌；确切的 `MeshNDInfo+0x38` 位→轴分配仍然是结构性的(参见下一个项目符号)。
- **`MeshNDInfo+0x38`.** popcount→dimension-count 的每比特分配是字节确认的；选择网格轴 0/1/2 的位是结构读数。中等的。
- **`MaybeMapShardIndexForHierarchicalSplit` (`0x138108e0`)。** `GetShardIndex` 后的每芯片核心分割重映射仅部分读取(`CoreIndex` / `LogicalDevicesPerChip` `idiv` at `0x13810932`)；没有追踪到完全复发及其无手术条件。低的。
- **阶段{零、一、二} DMA 调度。** `GetShardIndex` 被证明每环每相调用一次；没有跟踪每个阶段前进的轴以及 SFLAG 发送/接收相互作用(`GetRecvFlags`、`PipelinedAllGatherSlots`)——它们是 DMA 发射层。

---

## 交叉引用

- [集体概述](overview.md) — 基板分割和在此处路由 `all-gather` 的操作系列调度
- [SelectND策略](strategy-nd-picker.md) — 决定全聚集成为ND环的策略选择器； `[obj+0xa8]` 1-D-vs-ND 门
- [常量映射器](constant-mapper.md) — `CreateStaticNDRingReplicaInfoTable`：`InitDim` 读取的每轴 `device_id → ring-neighbor` 表(类型 0/1/2)
- [SPMD 链路计数成本](spmd-link-count-cost.md) — AllGather 成本分支(1D ÷2、2D ÷4)和 ICI `ResourceVector` 插槽
- [AllToAll 表](alltoall-tables.md) — A2A 障碍成员资格表生成(兄弟表系列)
- [ReduceScatter](reduce-scatter.md) — 在全归约分解中与全收集配对的归约分散阶段
- [降级轴摄取](degraded-axis.md) — 容错轴重新映射，将故障轴降级出主环
- [扭曲环面](../twist/overview.md) — 扭曲策略目标的非矩形环几何形状
- [路由](../routing/overview.md) — 消耗 `GetShardIndex` 偏移量的路由表/DMA 阶段层
- [ICI 面料](../ici/overview.md) — 芯片间互连 DMA 层
- [返回索引](../index.md)