张量分割/ND 平面

地址适用于 libtpu-0.0.40-cp314 轮中的 libtpu.so(build-id 89edbbe81c5b328a958fe628a9f2207d，build libtpu_lts_20260413_b_RC00；.text VMA == 文件偏移量 0xe63c000，.rodata VMA == 文件偏移 0x84a0000)。其他版本有所不同。

摘要

当嵌入类集合被卸载到 SparseCore(On-Pod 集体 §1.2 的 SC 卸载基底)时，基底必须先回答两个几何问题，然后才能构建环计划：有多少个 SparseCore 参与以及张量如何穿过它们，以及 集合的副本组实际投影到哪个圆环轴上。本页拥有这两个派生。第一个是设备分区：NumScOffloadDevices(参与 SC 计数)、tensor_split_factor(集体张量切割)以及 GetDimensionRings 内部的每轴环形设备分割。第二个是 ND 平面几何：ExtractNDPlaneInfo、它返回的 NDPlaneInfo / NDPlaneStrideInfo 结构体以及验证每个轴步幅的 IsNDPlaneSpanAcrossEntireDimension 投影。

ND 平面是密集 TensorCore ReplicaGroupsOnNDPlane 分解的 SparseCore 模拟 (SelectND策略 / 概述 §1.3)。其中密集路径存储 vector<MeshNDInfo> 并通过 ReplicaGroupsOnNDPlane(plane=2) 报告网格尺寸计数，SC 路径计算 NDPlaneInfo 描述符，其嵌入的 NDPlaneStrideInfo 携带每轴 optional<int32> 跨度步幅，并将尺寸计数报告为三个 has 位的 popcount — GetCollectiveNDPlaneDimensionCount。两者都计算副本组跨越的环面轴；它们仅在代表性方面有所不同。

这两个导数在下游相遇。 NumScOffloadDevices 总计通过 GetDimensionRings 在 X/Y/Z 环轴上进行乘法划分； NDPlaneInfo 尺寸计数门构建器的扭曲/ND 平面计数分支(cmp $3 3 轴扭曲门)。范围边界：所选逻辑颜色的“物理核心布局”位于 物理核心放置，“策略选择”(子平面、ND 环与扭曲环)位于 SelectND策略。此页面仅拥有分割因子分区、NDPlaneInfo 结构和每轴投影。

重新实现，合约为：

参与 SC 计数 NumScOffloadDevices = (TpuTopology SC-AvailableCoreCount / LogicalDevicesPerChip(SparseCore)) − num_embedding_devices，保留嵌入分区的卸载补充。
tensor_split_factor optional<int> 门：因子 >1 需要多个 SparseCore (!use_single_core)，唯一支持的非平凡因子是 2，它启用分割张量模式(颜色复制)。
每轴环形设备分割 segments = extent / devcount − 1，其中 devcount 作为跨轴的运行余数 — 卸载设备总数的乘法 ND 分区。
NDPlaneInfo / NDPlaneStrideInfo 字节布局和 IsNDPlaneSpanAcrossEntireDimension 投影 — 每轴步幅必须均匀划分环面尺寸范围才能使平面有效。


参与 SC 数量	`NumScOffloadDevices` `@0x1d6b8b00` → `long` (flag_utils.cc)
嵌入-保留兄弟	`NumEmbeddingDevices` `@0x1d6b8a00` → `long`(补分区)
张量分裂门	`ConstructConfigForCollectiveUniDirNDGroups<*>` `@0x133c82c0` / `0x133c2dc0` / `0x133cd800`(尾随 `optional<int>`)
每轴环分割	`GetDimensionRings` `@0x133df520` → `vector<RingConfigAttributes>` (`segments = extent/devcount − 1`)
ND平面提取器	`ExtractNDPlaneInfo` `@0x133bb940` → `StatusOr<NDPlaneInfo>` (offload_collective_config.cc)
每轴投影	`ExtractNDPlaneInfo::$_0` `@0x133bf700`(`IsNDPlaneSpanAcrossEntireDimension` RetCheck 三重)
尺寸计数	`GetCollectiveNDPlaneDimensionCount` `@0x133bb6e0` = `popcount(has_x + has_y + has_z)`
`NDPlaneInfo` 尺寸	~`0x28`(3 × `int32` + `NDPlaneStrideInfo` 在 `+0xc`)
`NDPlaneStrideInfo` 尺寸	`0x1c` (3 × `optional<int32>` + `bool across_cores_on_chip`)
圆环范围	`Target[+0x3b8][+0x58]` (X) / `+0x5c` (Y) / `+0x60` (Z) — 与密集拣选器 + 成本模型共享

1. SC卸载设备分区

设备分区以三个独立的部分回答“卸载的集合在多少个 SparseCore 上运行，以及张量如何在它们之间进行切割”：参与 SC 计数 (NumScOffloadDevices)、集合内张量切割 (tensor_split_factor) 和每轴环设备分割 (GetDimensionRings)。所有三个都是字节跟踪的。

1.1 `NumScOffloadDevices` — 参与 SC 数量

NumScOffloadDevices(ObjectView<TpuCompEnv>, const TpuTopology&) → long(@0x1d6b8b00，源 flag_utils.cc)计算可用于集体卸载的 SparseCore 逻辑设备 的数量。它是嵌入预留的补充：拓扑的总 SparseCores 除以每个芯片的 SparseCores 得出每个设备的 SC 计数；减去保留的 num_embedding_devices 就剩下可用卸载计数。

function NumScOffloadDevices(compEnv, topo):           // sub_1D6B8B00
    sc_total = topo[+0x94]                              // SC AvailableCoreCount (core type 2)
    ldpc_sc  = TpuTopology::LogicalDevicesPerChip(topo, 2)   // sub_20AD3020, SparseCore
    if ldpc_sc > 0:
        sc_dev = sc_total / ldpc_sc                     // signed idiv
    else:
        sc_dev = 0
    auto = AutoOr<long>::FromProtoOrDie(compEnv[+0x898])  // sub_1092F7E0 ("num_embedding_devices")
    n_emb = auto.engaged ? auto.value : 0               // dl bit0 = engaged
    CHECK(n_emb >= 0)                                   // "num_embedding_devices >= 0", line 1803 — FATAL
    CHECK(sc_dev >= n_emb)                               // "num_embedding_devices <= sc_per_device", line 1805
    return sc_dev - n_emb                               // = NumScOffloadDevices
```text

`TpuTopology[+0x94]` 字段是 SparseCore(核心类型 2)`AvailableCoreCount`。 `Target::CoresPerChip`(`@0x1d615b40`)读取的是每核型三重`TpuTopology[+0x7c + coretype·12]`的SC插槽(TC位于`+0x7c`，中间类型位于`+0x88`，SC位于`+0x94`)；它被`TpuTopology_MaybeAvailableSparseCoresPerLogicalDevice`(`@0xf6a1ea0`)独立确认为SC核心数，对于核心类型2称为`NumEmbeddingDevices`(`@0x1d6b8a00`)。

> **注意 —** 致命的 `CHECK` 源代码行在反编译中读取十进制 — `1803` 和 `1805` — 即 `0x70b` 和 `0x70d`。两个 CHECK 将结果括在有效范围内：`num_embedding_devices` 必须为非负数，并且不得超过 `sc_dev`(因此卸载补数不能为负数)。无论发生哪种违规，进程都会中止，并显示诊断 `"Invalid number of embedding devices specified"` (`.rodata @0x871a7c3`)。

同级 `NumEmbeddingDevices` (`@0x1d6b8a00`) 共享相同的 `sc_dev` 计算，但返回*保留*计数：`n_emb = engaged ? value : sc_dev`(当标志未设置时默认为所有 `sc_dev`)，有界 `0 <= n_emb <= sc_dev`(CHECK线路 `0x6f9` / `0x6fb`)。两个功能对`sc_dev`进行精确划分：`num_embedding_devices`(保留)+ `NumScOffloadDevices`(卸载)= `sc_dev`。

> **明白了 —** `NumScOffloadDevices` 是 SparseCore **逻辑设备** (`sc_total / ldpc_sc`) 的计数，而不是原始核心。从原始核心计数 `TpuTopology[+0x94]` 中减去 `num_embedding_devices` 的重新实现将在超级核心 SparseCore (`ldpc_sc == 2`) 上过度计数。除以 `LogicalDevicesPerChip(SparseCore)` 必须先于减法。

### 1.2 `tensor_split_factor` — 集体张量切割

模板化构建器 `ConstructConfigForCollectiveUniDirNDGroups<*>` 采用尾部 `optional<int>` `tensor_split` 作为其最后一个堆栈参数(`[rbp+0x10]`，复制到 `[rbp-0x108]` / `[rbp-0xf8]`)。来源 `offload_collective_config.cc`。门很小并且完全是错误驱动的：

```c
function ConstructConfigForCollectiveUniDirNDGroups(..., optional<int> tensor_split):
    split = tensor_split.has_value() ? tensor_split.value : 1   // cmovne, kDefaultTensorSplitFactor = 1
    if split >= 2:                                              // cmp eax, 2 ; jl skip
        RetCheck( !use_single_core.value_or(kDefaultUseSingleCore) )  // line 0x650
            // fail → "A larger than 1 tensor split factor requires more than one
            //         sparse core to split the tensor on."
        RetCheck( split == 2 )                                  // "tensor_split_factor.value_or(...) == 2", line 1558
            // fail → "We currently only support tensor split factor of 2 across two sparse cores."
        set tensor_split_mode = 2                               // [rbp-0x234] = 2 ; flag bit `or [..], 0x20`
        VLOG("Adopting split tensor mode.")                     // line 0x65c
        // mode-2 effect: "Twisted torus: duplicate colors as indicated by tensor split factor."

该因子通过 GenerationOption feed 降落在 OffloadConfig 原型字段 5 tensor_split_factor([variant+0x1c]，请参阅 SC-卸载配置生成器)。可选内容从公共 ND 包装器不变转发：

包装	地址	转发 `tensor_split`？
`ConstructConfigForAllReduceUniDirND`	`@0x133c2c80`	是 — `push [rbp+0x10]` `@0x133c2d01`
`ConstructConfigForReduceScatterUniDirND`	`@0x133ccbe0`	是 — 签名为 `optional<int>`
`ConstructConfigForAllGatherUniDirND`	`@0x133c76c0`	否 — 无 `optional<int>` 参数；接收默认的空可选

QUIRK — AllGather 无法在此版本中进行张量分割。其公共包装器 ConstructConfigForAllGatherUniDirND (@0x133c76c0) 没有 optional<int> 参数，因此其模板化构建器实例化始终看到默认值 1 因子，并且 split >= 2 门永远不会触发。只有 AllReduce 和 ReduceScatter 公开该旋钮。将 tensor_split 插入 AllGather 路径的重新实现是将二进制文件留下的旋钮连接起来。
明白了 — 唯一支持的重要因素正是 2。第二个 RetCheck("tensor_split_factor.value_or(kDefaultTensorSplitFactor) == 2"，源行 1558)拒绝 3, 4, … 与 "We currently only support tensor split factor of 2 across two sparse cores." 双核分割是密集 TensorCore 巨核数据分割的 SC 模拟； mode-2 的效果是重复颜色(“扭曲环面：由张量分割因子指示的重复颜色。”)，将张量分割到两个核心上，而不是将每个环的体积减半。此模式键控的每颜色发射归 物理核心放置 §5 (TensorSplitPerCoreClassifier) 所有。

1.3 `GetDimensionRings` — 每轴环形装置拆分

GetDimensionRings(const Target&, IciStrategyRingDim, int devcount, bool, bool megacore_aware) → vector<RingConfigAttributes> (@0x133df520) 是消耗运行设备计数并生成每轴环分解的分区器。 IciStrategyRingDim (1..7) 选择要读取的圆环轴范围； devcount 是沿该轴和以后的轴仍有多少设备需要分区。

function GetDimensionRings(target, ring_dim, devcount, b, megacore_aware):  // sub_133DF520
    validate ring_dim in 1..7                              // lea -7 ; cmp 0xfffffff9 ; jump table @0xae2eaac
    extent = target[+0x3b8][ X=0x58 | Y=0x5c | Z=0x60 ]    // chip torus dim extent
    set mesh/torus flag [-0x48] per switch arm
    ldpc_tc  = Target::LogicalDevicesPerChip(target, 0)    // sub_1D615B00 (TensorCore)
    megacore = (ldpc_tc >= 2) && megacore_aware            // setge AND'd with arg
    // THE SPLIT — how many ring segments this axis carries:
    segments = extent / devcount - 1                       // idiv ebx=devcount @0x133df670 ; dec eax @0x133df672 ; mov [-0xb8] @0x133df674
    ...                                                    // build RingConfigAttributes from segments + flags
```text

平面路径调用程序(AllGather 构建器，`@0x133c8d3e`..`@0x133c8d5c`)传递 `rcx = [rbp-0x150]` 中的**运行剩余设备计数**，从每轴双端队列元组的 `hi` 字段初始化。商 `extent / devcount` 成为下一个轴的 `devcount`，因此卸载设备在 X/Y/Z 环轴上**乘法**分区 — 每个轴消耗 `devcount` 设备，其余部分流向下一个轴。累积乘积累积在以 `[rbp-0x228]` 为根的 `std::__tree`(集)中——设备偏移/颜色索引图。在初始化该树之前，立即在 `rbx` 中获取 `NumScOffloadDevices`(`@0x133c89d3`..`@0x133c89e3`，通过 `GetTpuCompEnv` `@0x1d73de80`)，限制每轴环可能消耗的设备总数。

> **注意 —** `segments = extent / devcount − 1` 中的 `−1` 是将每个环的设备计数转换为*设备间跳数*的计数(N 个设备的环在包装之前有 N−1 个转发步骤)。它在`@0x133df670`–`@0x133df674`处进行字节确认(`idiv %ebx`然后`dec %eax`然后`mov %eax,-0xb8(%rbp)`；反编译`(int)v12 / v10 - 1`)，其中`v12`保存轴范围，`v10`保存设备计数。

---

## 2. `ExtractNDPlaneInfo` — ND 平面几何

`ExtractNDPlaneInfo` 将集合的副本组投影到物理环面上，并报告每个轴的*跨度步长* - 组沿该轴接触的芯片坐标之间的间距。这是密集 `ReplicaGroupsOnNDPlane` 分解的 SparseCore 对应项；它生成 `NDPlaneInfo` 描述符而不是 `vector<MeshNDInfo>`。

### 2.1 算法

```c
function ExtractNDPlaneInfo(target, device_assignment, hlo, Span<vector<int>> groups)  // sub_133BB940
        -> StatusOr<NDPlaneInfo>:                          // sret in rdi
    // (a) read chip torus extents (same offsets as the dense picker + cost model)
    X = target[+0x3b8][+0x58]; Y = target[+0x3b8][+0x5c]; Z = target[+0x3b8][+0x60]
    // (b) collect the distinct chip coordinate per axis
    for each member core_id in each replica group:
        loc  = FromGlobalCoreId(target, core_id)           // sub_133B7BC0 → TC core location
        coord = loc / LogicalDevicesPerChip                // idiv → chip-relative coordinate
        binary-search-insert coord into per-axis sorted-unique list   // memmove @0x133bbb30
    // (c) per axis (X, then Y, then Z) run the span lambda
    for axis in {X, Y, Z}:
        stride_slot = $_0(extent_axis, coord_set_axis, axis_name)      // sub_133BF700
    // (d) assemble the sret NDPlaneInfo: 3 ints + NDPlaneStrideInfo at +0xc
    return NDPlaneInfo{ size_x, size_y, size_z, NDPlaneStrideInfo{...} }

轴名称字符串("X" @0x8a106a1、"Y" @0x8a0f71b、"Z" @0x886531a)纯粹用于跨度检查失败时的诊断消息传递到 $_0 lambda。 duction-mod-LogicalDevicesPerChip 将每个全局核心 ID 折叠到其芯片坐标，因此同一芯片(兆核)上的两个核心映射到同一坐标，并向排序唯一集贡献一个条目。

2.2 `ExtractNDPlaneInfo::$_0` — 每轴投影 (`IsNDPlaneSpanAcrossEntireDimension`)

$_0 lambda (@0x133bf700) 计算一个轴的跨度步幅并验证该平面跨越整个环面尺寸。它写入 NDPlaneStrideInfo 槽：+8 处的 int32 大小和 +0xc 处的 bool has_size(相对于 lambda-this)，加上 +0 处的 StatusOr ok 标志。

function ExtractNDPlaneInfo::$_0(this, dim_size, coords, axis_name):   // sub_133BF700
    if coords.size() == 1:                                 // single coordinate ⇒ degenerate
        this[+8] = 0          // has_size = 0  (no span on this axis)
        this[+0xc] = 0
        this[+0]  = 1         // ok
        return
    stride = coords[1] - coords[0]                         // [-0x48]
    RetCheck( stride >= 1 )            // line 922; "Stride must be larger or equal to 1."
    RetCheck( stride < dim_size )      // line 923; "Stride must be less than the dimension size."
    RetCheck( dim_size % stride == 0 ) // line 925; "Stride must divide the dimension size."  ← idiv @0x133bf7b4
    // consistency: every adjacent pair must share the stride
    for k in 2 .. coords.size()-1:
        if coords[k] - coords[k-1] != stride:
            return FailedPrecondition(
                "All devices within a group must have the same stride along the "
                "dimension %s. Expected stride: %d but got %d.", axis_name, stride, observed)  // line 931
    this[+8]  = stride       // size = stride
    this[+0xc] = 1           // has_size = 1
    this[+0]  = 1            // ok
```text

这三个 RetCheck 是 **`IsNDPlaneSpanAcrossEntireDimension`** 门。关键的是第三个：`dim_size % stride == 0`。仅当平面的每轴步幅均匀地划分圆环尺寸范围时，平面才是有效的 ND 平面 - 即跨步坐标集干净地围绕整个圆环轴。不划分范围的步幅会留下部分环，SC 卸载基板不会发出该部分环。

> **明白了 —** lambda **不**仅仅读取 `coords[1] - coords[0]`；然后，它“验证”排序坐标集中的每个相邻对都具有相同的步幅(源代码行 931，错误 `"All devices within a group must have the same stride along the dimension %s."`)。成员沿轴不均匀间隔的组(有效的 HLO 副本组，但不是干净的环面平面)在这里被拒绝，而不是由第一对默默近似。原始单对读数是必要的，但还不够；完整的每对扫描才是真正的合约。
>
> **注意 —** 源代码行根据原始结果确认字节精确：`stride >= 1` at `922` (`0x39a`)、`stride < dim_size` at `923` (`0x39b`)、`dim_size % stride == 0`在`925` (`0x39d`)。第二次检查的面向用户的字符串是 `"Stride must be less than the dimension size."`，第三次检查的面向用户的字符串是 `"Stride must divide the dimension size."`

### 2.3 `NDPlaneInfo` / `NDPlaneStrideInfo` 结构布局

这两个结构体由其 `ToString` 方法(`NDPlaneInfo::ToString` `@0x10fdf2a0`、`NDPlaneStrideInfo::ToString` `@0x10fe62e0`)固定，并由两个使用者 `GetCollectiveNDPlaneDimensionCount` 中的字段读取独立交叉检查(`@0x133bb6e0`)和`GetMinorToMajorOrder`(`@0x133c1c40`)。

| 结构 | 偏移 | 型 | 字段 |
|---|---|---|---|
| `NDPlaneInfo` | `+0x00` | `int32` | `size_x`(平面跨度，X轴) |
|  | `+0x04` | `int32` | `size_y` |
|  | `+0x08` | `int32` | `size_z`(也是迭代/平面边界 `GetMinorToMajorOrder` 在 `+0x8` 处读取) |
|  | `+0x0c` | `NDPlaneStrideInfo` | 嵌入步幅描述符 |
| `NDPlaneStrideInfo` | `+0x00` | `int32` | `stride_x`(`optional<int32>`值) |
| (位于 `NDPlaneInfo+0xc`) | `+0x04` | `bool` | `has_stride_x` |
|  | `+0x08` | `int32` | `stride_y` |
|  | `+0x0c` | `bool` | `has_stride_y` |
|  | `+0x10` | `int32` | `stride_z` |
|  | `+0x14` | `bool` | `has_stride_z` |
|  | `+0x18` | `bool` | `across_cores_on_chip`(兆核/双核标志) |

`sizeof(NDPlaneStrideInfo)` 为 `0x1c`； `sizeof(NDPlaneInfo)` 是 `~0x28`(三个整数 + 嵌入的 `0x1c` 字节步长信息，向上舍入)。相对于父级 `NDPlaneInfo`，嵌入的步幅字段位于：`stride_x`@`+0xc` / `has_x`@`+0x10`、`stride_y`@`+0x14` / `has_y`@`+0x18`、`stride_z`@`+0x1c` / `has_z`@`+0x20`、`across_cores`@`+0x24`。

`NDPlaneInfo::ToString` 标签 (`@0x10fdf2a0`) 是字节精确的 `"size_x: "`、`"size_y: "`、`"size_z: "`，嵌入式 `NDPlaneStrideInfo` 的尾随 `"stride_info: "`； `NDPlaneStrideInfo::ToString`将其三个字段标记为`"stride_x: "`、`"stride_y: "`、`"stride_z: "`。因此，三个顶级整数是**平面跨度大小** `(size_x, size_y, size_z)`，并且嵌入的结构携带每轴**步幅**。

> **QUIRK —** 尺寸/步幅配对是反直觉的部分。顶级 `NDPlaneInfo` 整数是平面*跨度* (`size_*`)，而 `$_0` lambda 计算的*跨度*(`coords[1] - coords[0]`) 存储在嵌入的 `NDPlaneStrideInfo` 中作为 `stride_*`。将两者混为一谈的重新实现者(将步幅写入顶级 int )将错误地输入 `GetMinorToMajorOrder`，它读取嵌入的 `stride_*`/`has_*` (`@0x133c1c40`: `has_x`@`+0x10`， `size_x`@`+0xc`(相对于`NDPlaneInfo`)，不是顶级的`size_*`。

### 2.4 ND平面尺寸计数

`GetCollectiveNDPlaneDimensionCount` (`@0x133bb6e0`) 是密集 `ReplicaGroupsOnNDPlane(plane=2).num_mesh_dims` 的 SC 模拟。它调用 `ExtractNDPlaneInfo` 并将结果减少为单个整数：副本组跨越的环面轴的数量。

```c
function GetCollectiveNDPlaneDimensionCount(target, device_assignment, collective):  // sub_133BB6E0
    CHECK(collective != nullptr)                           // line 845
    groups = GetPhysicalDeviceGroups(collective, device_assignment)
    info   = ExtractNDPlaneInfo(target, device_assignment, collective, groups)
    if !info.ok: return info.status                        // AddSourceLocation line 852 / 849
    return info.has_stride_x + info.has_stride_y + info.has_stride_z   // sum of the 3 has-bytes

维度计数实际上是三个 has_* 字节的总和 - 在 @0x133bb772..@0x133bb786 进行字节确认(反编译行 *((_DWORD *)this + 2) = v26 + v27 + v28;，其中 v26/v27/v28 是从本地读取的三个 has-bits) NDPlaneInfo 位于 +0x18/+0x20/+0x28)。因为每个 has_* 都是 0 或 1，所以总和就是三个轴上的 popcount。该计数是构建者的 cmp $3 扭曲门(3 暗淡 → k_2k_2k / k_k_2k 扭曲)和 AllGather/ReduceScatter ND 平面计数门(1..4 暗淡)测试的结果。

注意 — 密集路径和 SC 路径通过不同的路由计算相同的数量。密集的 GetCommunicationMultiplier (@0x127a16c0, 概述 §3) 返回 ReplicaGroupsOnNDPlane(plane=2).num_mesh_dims + 1 作为链接计数除数； SC GetCollectiveNDPlaneDimensionCount 返回不含 +1 的裸 popcount。 +1 是密集侧的成本模型约定，而不是基础轴计数的差异 - 两者都对副本组跨越的环面轴进行计数。

2.5密集`ReplicaGroupsOnNDPlane`投影(`MeshNDInfo`构建器)

§2.1 的 SC ExtractNDPlaneInfo 的密集 (TensorCore) 对应项是 ReplicaGroupsOnNDPlane (@0x1c890960)。它是AllGather ND 环和ReduceScatter 选择器调用的构建器，用于“决定”维数和“构建”每轴vector<MeshNDInfo>。输入函数本身不包含不坐标数学：它(a)将设备分配呈现为字符串并序列化拓扑(TpuTopologySerdes::Distill→TpuTopologyArgs::ToProto→tsl::SerializeToStringDeterministic)，(b)采用nd_plane_cache_mutex并在NDPlaneCacheKey → optional<vector<MeshNDInfo>>缓存中查找组合键(@0x225799b8，受保护的单例 GetNDPlaneCache)，并且 (c) 在未命中时调用每轴 lambda ReplicaGroupsOnNDPlaneImpl::$_0 (@0x1c896400) 每个网格轴一次 - axis ∈ {0,1,2} 作为 0x100000000 | axis 传递 - 然后将生成的 optional<vector<MeshNDInfo>> 存储到缓存中。

function ReplicaGroupsOnNDPlane(target, device_assignment, device_list, n_dim, b):  // sub_1C890960
    CHECK(device_assignment != nullptr)
    key.dev   = render(device_assignment)                  // "<id>_<d0>,<d1>,…"
    key.topo  = SerializeToStringDeterministic(Distill(target.topology).ToProto())
    key.n_dim = n_dim; key.b = b
    lock(nd_plane_cache_mutex)
    CHECK(n_dim == 1 || n_dim == 2 || n_dim == 3)          // line 1155, group_utils.cc
                                                           // "…only supports dimension n_dim = 1, 2 or 3."
    if !cache.contains(key):
        v = nullopt
        for axis in {0, 1, 2}:                             // built per mesh axis
            v = ReplicaGroupsOnNDPlaneImpl::$_0(&state, groups, 0x100000000 | axis)   // sub_1C896400
            if !v.has_value: break                         // projection failed ⇒ no plane
        cache[key] = v
    result = cache[key]                                    // optional<vector<MeshNDInfo>>
    unlock(nd_plane_cache_mutex)
    return result
```text

lambda `$_0` (`@0x1c896400`) 是投影实际发生的地方；它是 `ExtractNDPlaneInfo::$_0` (§2.2) 的密集模拟，但它发出 `MeshNDInfo` 环形几何形状，而不是单个跨度步长。

```c
function ReplicaGroupsOnNDPlaneImpl::$_0(state, replica_groups, axis_flag):  // sub_1C896400
    n_dim = state.n_dim                                    // ***(int***) = 1 | 2 | 3
    if n_dim == 1:
        // trivial 1-D plane: one MeshNDInfo whose single axis lists all group members
        return optional(vector<MeshNDInfo>{ MeshNDInfo_1D(replica_groups) })
    ldpc = LogicalDevicesPerChip(target, /*TensorCore*/0)  // megacore ⇒ 2
    // mesh extents seeded from the chip torus (X,Y,Z at target[+0x58]/[+0x60]/[+0x161,…])
    if axis_flag & 0x100000000:                            // a specific axis was requested
        a = axis_flag & 3                                  // 0|1|2
        mesh_extent[a] *= ldpc                              // fold the per-chip sub-cores into this axis
    // multi-slice fan-out
    num_slices = GetMultiSliceTopology(target) ? GetNumSlices(target) : 1
    alloc PerSliceReplicaData[num_slices]                  // 88-byte stride per slice
    for each group:
        for each device d in group:
            (core_id, slice) = GetMegascalePerSliceCoreIdAndSliceId(target, da, d)   // sub_1C8906E0
            loc   = TensorCoreLocationForLogicalDeviceId(target, da, core_id, nullopt) // sub_1C8904E0
            (cx,cy,cz) = loc.chip_coordinates()
            idx = mixed_radix_linearize((cx,cy,cz), mesh_extent)    // Horner over mesh strides
            per_slice[slice].grid[idx] = group_member_ordinal
    // per (slice, group) emit one MeshNDInfo by dispatching on n_dim
    for each slice, each group:
        if n_dim == 2: m = ReplicaGroupForm2DRing(group, …, mesh_extent)   // sub_1C88E6E0
        if n_dim == 3: m = ReplicaGroupsOn3DPlane(group, …, NDTopologyInfo) // sub_1C8901E0
        if !m.has_value: return nullopt                    // group does not fit an n_dim plane
        out.push_back(m)
    return optional(out)

重新实现者必须保留的三个事实，全部经过反编译验证：

每轴子核心折叠。 当请求特定轴 (axis_flag & 0x100000000) 时，该轴的网格范围将预先乘以 LogicalDevicesPerChip(mesh_extent[axis] *= ldpc，v253[8*(axis&3)+16] *= v51 存储)。在超级核 TensorCore (ldpc == 2) 上，所选环轴携带两个片上内核；另外两个轴停留在芯片粒度上。在所有三个轴上针对原始芯片范围进行线性化的重新实现将错误地放置超级核心的第二个核心。
chip_coordinates，不是原始核心 ID。 每个设备都通过 TensorCoreLocationForLogicalDeviceId → chip_coordinates 进行投影，然后针对(折叠)网格范围进行混合基数线性化 - 与霍纳式 coord + extent·(…) 重复相同，SC 路径减少 mod LogicalDevicesPerChip。投影是芯片坐标放置，而不是核心 ID 排序。
n_dim 调度 + 短路。 n_dim == 1 返回一个简单的单轴 MeshNDInfo，没有任何环结构； n_dim == 2将各组调度至ReplicaGroupForm2DRing(@0x1c88e6e0)； n_dim == 3 至 ReplicaGroupsOn3DPlane (@0x1c8901e0)。任何不适合 n_dim 轴平面的组都会使帮助器返回 nullopt，这会中止整个投影(v.has_value 在入口循环中中断)，并且正是 AllGather/ReduceScatter 选择器测试的“设备列表未投影到 k 轴平面上”信号。

注意 — 这解决了 allgather-nd-ring.md“未解决的问题”条目，该条目将投影数学放置在 0x1c891402 内部 ReplicaGroupsOnNDPlane 处。入口函数(0x1c890960)仅构建缓存键并进行调度；实际的坐标投影体是 0x1c896400 处的每轴 lambda ReplicaGroupsOnNDPlaneImpl::$_0。 n_dim == 1 || 2 || 3 断言("…only supports dimension n_dim = 1, 2 or 3."，group_utils.cc 第 1155 行)是入口函数的唯一内联检查。
明白了 — 多切片路径不是可选的簿记。设置 GetMultiSliceTopology 时，lambda 为每个 GetNumSlices 分配一个 PerSliceReplicaData(88 字节步长)，并且由 GetMegascalePerSliceCoreIdAndSliceId (@0x1c8906e0) 选择设备的插槽，这会返回*每片核心 id 和片索引。忽略切片索引的单切片重新实现会将来自不同切片的设备碰撞到多切片(超大规模)拓扑上的同一网格单元中。

3. 两个推导如何满足

当卸载配置构建器布置每色环时间表时，设备分区(§1)和 ND 平面几何形状(§2)一起使用。该表总结了数据流。

数量	源码	角色
SC 可用CoreCount	`TpuTopology[+0x94]`(核心类型2)	拓扑中 SparseCore 总数
`LogicalDevicesPerChip(SparseCore)`	`TpuTopology::LogicalDevicesPerChip(2)` `@0x20ad3020`	每个芯片稀疏核心数(兆核 ⇒ 2)
`sc_dev = total / ldpc_sc`	`NumScOffloadDevices` / `NumEmbeddingDevices`	SC 逻辑设备计数
`num_embedding_devices`	`AutoOr<long>` 的 `compEnv[+0x898]`	预留嵌入分区
`NumScOffloadDevices`	`sc_dev − num_embedding_devices`	卸载可用 SC 计数(构建器 `rbx`)
`tensor_split_factor`	`optional<int>` 构建器 arg → 原型字段 5	集体张量切割(`==2` ⇒ 分割模式)
每轴环段	`GetDimensionRings`：`extent / devcount − 1`	环是如何沿每个轴切割的
ND平面尺寸计数	`popcount(has_x + has_y + has_z)`	环面轴集体跨度(扭转门)
每轴跨距	`$_0` 拉姆达：`coords[1] − coords[0]`、`stride \| extent`	`NDPlaneStrideInfo` `stride_x/y/z`
短轴到长轴顺序	`GetMinorToMajorOrder` `@0x133c1c40`(读取为 `NDPlaneStrideInfo` + `topo[+0xa3]`)	跨跨轴环遍历顺序

NumScOffloadDevices 总数限制了 GetDimensionRings 跨 X/Y/Z 分区的正在运行的 devcount； NDPlaneInfo 尺寸计数选择构建器的扭曲/ND 平面分支； tensor_split_factor 修改每个颜色的发射(模式 2 中的颜色重复)。然后，GetMinorToMajorOrder 将每轴 NDPlaneStrideInfo 跨度大小转换为短轴到长轴，排序每色环迭代，由 topo[+0x3b8][+0xa3] across_cores_on_chip 标志门控(如果 3D 平面缺少轴，则在第 2295/2297/2299 行出现致命 "Stride x/y/z should be set for 3D plane."跨步)。

4. 验证说明

针对 libtpu.so v0.0.40 (build-id 89edbbe8…) 的 IDA 反编译进行交叉检查：
NumScOffloadDevices (@0x1d6b8b00) — sc_total = topo[+0x94]； ldpc_sc = LogicalDevicesPerChip(2)； sc_dev = sc_total/ldpc_sc(idiv、ldpc_sc>0 防护)； AutoOr<long>::FromProtoOrDie(compEnv[+0x898]) 接合位；两个致命检查("num_embedding_devices >= 0" 第 1803 行、"num_embedding_devices <= sc_per_device" 第 1805 行)； return sc_dev − n_emb — 所有字节精确。兄弟 NumEmbeddingDevices (@0x1d6b8a00) 确认为补充分区。
tensor_split_factor 门(AllGather 构建器 @0x133c82c0)— value_or(1)； cmp 2 >= 2 门；重新检查 "!use_single_core.value_or(kDefaultUseSingleCore)" → "A larger than 1 tensor split factor requires more than one sparse core to split the tensor on."； RetCheck "tensor_split_factor.value_or(kDefaultTensorSplitFactor) == 2"(第1558行)→"We currently only support tensor split factor of 2 across two sparse cores."；视频博客"Adopting split tensor mode."；模式 2 效果 VLOG "Twisted torus: duplicate colors as indicated by tensor split factor." — 所有字节精确。 AR(push [rbp+0x10] @0x133c2d01)/RS 转发已确认； AllGather 包装器 @0x133c76c0 确认没有 optional<int> 参数。
GetDimensionRings (@0x133df520) — X/Y/Z 范围 [+0x3b8][0x58/0x5c/0x60]； LogicalDevicesPerChip(0)(TC)巨核检测；在 @0x133df670 (idiv %ebx) / @0x133df672 (dec %eax) / @0x133df674 (mov %eax,-0xb8(%rbp); 反编译 (int)v12 / v10 − 1) 处字节确认的拆分 extent / devcount − 1 — 准确。
ExtractNDPlaneInfo::$_0 (@0x133bf700) — coords.size()==1 快速路径 (has_size=0)； stride = coords[1] − coords[0]；第 922/923/925 行的三个 RetCheck("Stride must be larger or equal to 1." / "Stride must be less than the dimension size." / "Stride must divide the dimension size.")；每对步幅一致性扫描(第 931 行)； has_size=1, size=stride 成功存储 — 所有字节都是精确的。
NDPlaneInfo / NDPlaneStrideInfo 布局 — NDPlaneInfo::ToString @0x10fdf2a0 标签 "size_x: "/"size_y: "/"size_z: "/"stride_info: "； NDPlaneStrideInfo::ToString @0x10fe62e0 标签 "stride_x: "/"stride_y: "/"stride_z: "/"across_cores_on_chip: "，读取位置为 +0/+4、+8/+0xc、 +0x10/+0x14、+0x18；两个消费者GetCollectiveNDPlaneDimensionCount @0x133bb6e0(has位)和GetMinorToMajorOrder @0x133c1c40(has@+0x10/+0x18/+0x20， size@+0xc/+0x14/+0x1c) 同意 ToString 派生布局 — 字节精确。
GetCollectiveNDPlaneDimensionCount (@0x133bb6e0) — *((_DWORD *)this + 2) = v26 + v27 + v28(三个 has 字节的总和)已确认；检查 "collective != nullptr" 第 845 行； AddSourceLocation 线 849/852 — 准确。
密集 ReplicaGroupsOnNDPlane (@0x1c890960) + lambda ReplicaGroupsOnNDPlaneImpl::$_0 (@0x1c896400) — 入口函数构建 NDPlaneCacheKey(设备分配渲染 + TpuTopologySerdes::Distill/ToProto/SerializeToStringDeterministic + n_dim + bool) 在 nd_plane_cache_mutex 下，断言 n_dim == 1||2||3("…only supports dimension n_dim = 1, 2 or 3."，group_utils.cc 第 1155 行)，并且在缓存未命中时，每个轴调用一次 lambda 0x100000000 | axis (axis ∈ {0,1,2}) — 字节确认。 lambda的每轴子核心折叠mesh_extent[axis&3] *= LogicalDevicesPerChip(0)(v253[8*(a5&3)+16] *= v51)，n_dim==1琐碎-MeshNDInfo短路(**(int**)a2 == 1臂)，每设备TensorCoreLocationForLogicalDeviceId(@0x1c8904e0)→ chip_coordinates → 混合基数线性化，通过 GetMultiSliceTopology/GetNumSlices/GetMegascalePerSliceCoreIdAndSliceId (@0x1c8906e0) 进行多片扇出，具有 88 字节 PerSliceReplicaData 步长，并将 n_dim 调度到ReplicaGroupForm2DRing (@0x1c88e6e0) / ReplicaGroupsOn3DPlane (@0x1c8901e0) 和 nullopt-on-no-fit — 全部经过反编译验证。
[低] 由结构/标签确认，而不是由独立的数字消费者确认：
顶级 NDPlaneInfo size_x/y/z 整数携带的精确算术(平面范围与 沿轴的设备计数)：标签是字节读取的，并且 size_z (+0x8) 用作 GetMinorToMajorOrder 绑定的迭代/平面，但是ExtractNDPlaneInfo 的 sret 尾部 (@0x133bc674..@0x133bc880) 中的生产者算术未单独解码 - 大小语义是从 "size_x: " 标签和消费者使用中推断出来的。
NDPlaneStrideInfo (SC) 与密集 MeshNDInfo (TC) 的统一 - 两者都报告每轴跨度，但逐场对应关系并未关闭； SC维度计数使用has位的popcount，TC使用ReplicaGroupsOnNDPlane.num_mesh_dims。

交叉引用

范围边界(本页的邻居)

物理核心放置 — 所选逻辑颜色落在具体的物理 SC 核心上，tensor_split_mode==2 每核心发射由 TensorSplitPerCoreClassifier 键入
SelectND策略 — 致密基底 ND 策略选择(子平面/ND 环/扭曲/跨步)，TC ReplicaGroupsOnNDPlane 模拟本页的 NDPlaneInfo 镜子

SparseCore 卸载基板

SC-卸载配置生成器 — 携带 tensor_split_factor(字段 5)的 *OffloadConfig 原型，以及此页面的门位于内部的 ConstructConfigForCollectiveUniDirNDGroups<*> 构建器
分层种类 — AutoOr<bool> 扁平与分层拆分卸载构建器调度
SC 核心选择(卸载) — SparseCoreConfig.offload op 类型分类 + 核心选择

基板图 + 兄弟子系统

On-Pod 集体 — 剖面图 — 衬底分割、SC 卸载门和共享物理环面网格分解
AllToAll 表 — 全部到全部/参差不齐的链接表 (EstimatePhysicalLinksUsed)
返回索引

张量分割/ND 平面 ​

摘要 ​

1. SC卸载设备分区 ​

1.1 NumScOffloadDevices — 参与 SC 数量 ​

1.3 GetDimensionRings — 每轴环形装置拆分 ​

2.2 ExtractNDPlaneInfo::$_0 — 每轴投影 (IsNDPlaneSpanAcrossEntireDimension) ​

2.5密集ReplicaGroupsOnNDPlane投影(MeshNDInfo构建器) ​

3. 两个推导如何满足 ​

4. 验证说明 ​

交叉引用 ​

范围边界(本页的邻居) ​

SparseCore 卸载基板 ​

基板图 + 兄弟子系统 ​