Embedding Minibatching 分解

本页上的每个地址、opcode、operand 索引、custom-call 名称和源码行标签都逐字节读取自 libtpu-0.0.40-cp314 wheel 中的 libtpu.so（build-id 89edbbe81c5b328a958fe628a9f2207d，build libtpu_lts_20260413_b_RC00；未 strip，nm -C 可解析每个方法）。.text VMA 等于其文件偏移（0xe63c000）；.rodata 位于 0x84a0000；.data.rel.ro 是 VMA − 0x200000（operand-name 表，加载时由 R_X86_64_RELATIVE 填充）。地址适用于此 build；其他版本会有差异。

摘要

本页是位于 SC scan lowering 之上的 HLO 层分解：前端的整体稀疏嵌入矩阵乘 custom-call（SparseDenseMatmulWithMinibatchingOp）如何被拆分为每个 minibatch、每个物理 SparseCore core 一个内部 SparseDenseMatmulOp，每个 split 的 per-core CSR 窗口如何寻址，以及结果如何馈入下游 SampleCombiner emitter 和 valency loop 所消费的 SC-dialect op 集。它由 xla::tpu::sparse_core 中三个同级 HLO/MLIR pass 负责：MinibatchingDecomposition（per-minibatch CSR-slice 算术 + 前向/反向 pass 名册）、EmbeddingDataFormattingDecomposer（dense per-table ↔ SC-packed-stacked-table activations/gradients 适配器）和 PackedOperandsLowering（在 packed-width operands 上重新创建 dialect sparse_core::SegmentedScanOp 的 MLIR full-conversion pass）。

唯一最关键、也最容易被误读的结构事实是：名为 DynamicSliceCsr 的 op 不是 HLO kDynamicSlice（0x36）。CreateDynamicSliceCsr（0x13489ea0）完全不会发出 dynamic-slice opcode。它用 GetCoreIndex custom-call、GetPaddedRowCount granule clamp，以及纯整数乘/加链（三个 kMultiply=0x4b 加一个 kAdd=0x3）构造一个 {sliced-csr, base-offset, padded-count} 三元组。实际的 per-minibatch CSR row-pointer 窗口由前向 pass 随后发出的内部 SparseDenseMatmulOp custom-call 按该元组参数化后切片。DynamicSliceCsr/GetCoreIndex 是 SparseCoreOperationType custom-call 名称（op-types 0x10 / 0xc），不是 HLO opcode。

第二个关键事实是三个 pass 之间的职责分工。MinibatchingDecomposition 是唯一触碰 CSR row-pointers 的 pass；它是 SegmentedScan 最终归约的 segment-id 来源。EmbeddingDataFormattingDecomposer 在 dense per-table XLA layout 与 SC packed stacked-table layout 之间重排activations 和 gradients；它不是 CSR/id/gain operand 重排器。PackedOperandsLowering 是 dialect 层 op-packing pass，它通过 unpack→create→pack rewrite 在目标 packed width 上重新创建 SegmentedScanOp（以及约 40 个其他 AluEp op）。重新实现者必须把这三件事分开。

本页由三个单元加下游绑定组成：minibatching decomposition pipeline（op 识别、CreateDynamicSliceCsr、granule clamp、前向/反向名册、while-fusion vs no-while）、operand partition（per-core CSR base-offset 算术 + 前向/反向 operand-name 表）、packed-operands lowering（SegmentedScanOp 重建器），以及绑定到 DotCombiner emitter 产生的 gather→sort→uniquify→reduce→scatter SC Stream-op DAG。

Decomposition pass	MinibatchingDecomposition::RunImpl（0x1348f940）— 按 custom-call 名称扫描 op，构造 ArgSpecs，分派前向/反向
Op recogniser	minibatching_decomposer_util::IsSparseDenseMatmulWithMinibatchingOp（0x13c86da0）— IsCustomCall("…WithMinibatchingOp", 35) OR IsCustomCall("…GradOptimizerUpdateWithMinibatchingOp", 54)
Per-minibatch slice	MinibatchingDecomposition::CreateDynamicSliceCsr（0x13489ea0，0x8c0 B）→ {sliced-csr, base, padded} 三元组。无 HLO kDynamicSlice。
Padded count	sparse_dense_matmul_decomposer_util::GetPaddedRowCount（0x13c90280）= max( max(GranuleBytes/4, num), cfg[+0x948]→[+0x94] )，由 SupportsSparseCore() gate
Custom-call op-types	GetCoreIndex = 0xc，DynamicSliceCsr = 0x10（SparseCoreOperationTypeToString 0x14b7f480）
发出的 HLO opcode	kMultiply = 0x4b（×3），kAdd = 0x3（×1）— 已对照 StringToHloOpcode init（0x1e5ef040）验证
分解后的内部 op	GetSparseDenseMatmulOpCustomCallTarget（0x13c86e60）→ "SparseDenseMatmulOp"（19 B）
Data-format adapter	EmbeddingDataFormattingDecomposer::RunImpl（0x1368b4a0）— op-types 0x1a..0x1d；Sc/Tc 由 EnableEmbeddingDataFormattingOffload gate
Packed-op lowering	ScanOpLowering<SegmentedScanOp>::matchAndRewrite（0x135f3000）— unpack→getReductionOp→SegmentedScanOp::create→pack
下游 emitter	SparseDenseMatmulDotCombinerEmitter::Emit（0x1332bda0）→ EmitValencyLoop / EmitVectorizedLoop / EmitSampleCombiner

---

单元 1 — Minibatching Decomposition Pipeline

minibatching 存在的原因

TPU 嵌入查找会在每个训练 step 产生一个拼接的 CSR（compressed-sparse-row）结构，描述每个样本触碰了哪些嵌入表行。当全局 lookup 过大，无法一次处理完时，前端会将 custom-call 标记为 minibatching op：该 lookup 将被拆分为多个 sub-batch 独立处理，再重组结果。MinibatchingDecomposition 是 HLO pass，它将这个带标签的 custom-call 转换为每个 (minibatch, physical SC core) 对一个具体的内部 SparseDenseMatmulOp；这个 per-pair op 正是 SC scan datapath 实际为其发出 sequencer program 的单位。

结构上重要的一点是，该分解纯粹是 HLO graph rewrite。它不发出 SC dialect、不产生 SPMEM 流量、不产生 sequencer 指令；它产生一棵由 HLO custom-call 与算术组成的树，其叶子是 per-minibatch 内部 op。SC-dialect lowering 发生在之后，当每个内部 SparseDenseMatmulOp 自身被降低时（单元 4）。

Op 识别 — 按 custom-call 名称

RunImpl（0x1348f940）遍历 module，并按 custom-call target string 识别目标 op，而不是按 opcode 或 attribute。识别器 IsSparseDenseMatmulWithMinibatchingOp（0x13c86da0）是两个名称的析取，已从反编译函数体逐字节确认：

// IsSparseDenseMatmulWithMinibatchingOp (0x13c86da0) — decompile-exact
bool IsSparseDenseMatmulWithMinibatchingOp(const HloInstruction *op) {
    if (op->IsCustomCall("SparseDenseMatmulWithMinibatchingOp", 35))           // forward
        return true;
    return op->IsCustomCall("SparseDenseMatmulGradOptimizerUpdateWithMinibatchingOp", 54);  // backward
}
```text

匹配后，`GetArgSpec`（`0x13c87040`）选择每个 op 的 operand spec：`ForwardPassArgSpec`（vtable `0x21937cc8`）或 `BackwardPassArgSpec`（vtable `0x219382c0`）；`GetSparseDenseMatmulOpCustomCallTarget`（`0x13c86e60`）把匹配到的 *minibatching* op 映射到**分解后的内部 op 名称**。映射来自反编译字符串赋值分支：

| 匹配的 custom-call（长度 B） | 分解后的内部 op |
|---|---|
| `SparseDenseMatmulWithMinibatchingOp`（35） | `SparseDenseMatmulOp`（19） |
| `SparseDenseMatmulGradOptimizerUpdateWithMinibatchingOp`（54） | grad-with-optimizer-update inner（`kSparseDenseMatmulGradOpWithOptimizerUpdate`） |
| `SparseDenseMatmulCustomCombinerMegachipOp`（41） | `SparseDenseMatmulCustomCombinerOp` |
| `…CustomCombinerTcCombinerMegachipOp` | `…CustomCombinerTcCombinerOp` |

> **注意 — 识别是字符串键控的，因此 op 词汇表是一个闭合的 `.rodata` 字符串集合。** 匹配是 `HloInstruction::IsCustomCall(target_string, length)` 对 literal byte strings，而不是 enum 比较。重新实现者应通过字符串匹配 custom-call target 来复现；megachip/custom-combiner 变体是同一识别器家族中的额外 literal 字符串。

### `CreateDynamicSliceCsr` — per-minibatch slice 描述符（不是 dynamic-slice）

`CreateDynamicSliceCsr`（`0x13489ea0`）是 partition 的核心。它的名称来自它产生的 *SC custom-call*，而不是任何 HLO dynamic-slice opcode。它的工作是：为一个 minibatch 构造 `{sliced-csr, base-offset, padded-count}` 三元组，内部 `SparseDenseMatmulOp` 将用它精确读取拼接 CSR row-pointers 的切片。反编译的 HLO-builder 调用序列按发出顺序如下：

```text
CreateDynamicSliceCsr (0x13489ea0)  — decompile-confirmed op DAG, in order
  guard  RetCheckFail "max_ids_per_partition > 0"   (minibatching_decomposer.cc:154) if num <= 0
  0  padded = GetPaddedRowCount(target, num)         // 0x13c90280  — granule clamp (see below)
  1  C_pad  = CreateConstant( LiteralUtil::CreateR0<int>(padded) )   // s32 scalar
  2  GCI    = CreateCustomCall(s32[1], "GetCoreIndex", {})           // op-type 0xc; runtime per-core index
  3  mul1   = CreateBinary(s32[1], MULTIPLY/*0x4b*/, a6,    C_pad)   // a6 * padded
  4  mul2   = CreateBinary(s32[1], MULTIPLY/*0x4b*/, GCI,   mul1)    // GCI * a6 * padded
  5  mul3   = CreateBinary(s32[1], MULTIPLY/*0x4b*/, C_pad, a7)      // padded * a7
  6  base   = CreateBinary(s32[1], ADD/*0x3*/,       mul2,  mul3)    // per-core CSR base offset = padded·(GCI·a6 + a7)
  7  dsc    = CreateCustomCall(…, "DynamicSliceCsr", {csr, base, C_pad})  // op-type 0x10; 3 operands
  8  gte    = CreateGetTupleElement(shape, dsc, 0)
  9  tuple  = CreateTuple({ gte, base, C_pad })       // StatusOr<HloInstruction*> 3-tuple

这三个乘法都是 opcode 0x4b（MULTIPLY），单个加法是 0x3（ADD）；二者均对照按字母排序的 StringToHloOpcode init 表（0x1e5ef040）读取，其中 add=0x3、multiply=0x4b、dynamic-slice=0x36。这里从未发出 0x36，此函数中没有任何 HLO dynamic-slice。shape 通过 ShapeUtil::MakeValidatedShape(S32 /*PrimitiveType 4*/, …) 构造。

注意 — CreateDynamicSliceCsr 不发出 HLO dynamic-slice。 反编译函数体只发出 kMultiply（×3）和 kAdd（×1） 作为 HLO opcode，另有 Constant/CustomCall/GetTupleElement/Tuple。"DynamicSliceCsr" 字符串是 SparseCoreOperationType custom-call 名称（op-type 0x10），不是 HloOpcode::kDynamicSlice。per-minibatch CSR row-pointer 窗口在内部 SparseDenseMatmulOp 内部切片，由该元组携带的 {base, padded} 参数化。

注意 — early-return guard 字符串。 num <= 0 guard 通过消息 "max_ids_per_partition > 0"（platforms/xla/sparse_core/hlo/minibatching_decomposer.cc:154）执行 RetCheck，第二个关于 inner-call result 的 RetCheck 检查 "concatenated_csr_tuple.size() > kCsrTupleRowPtrIndex"（line 177）。两个字符串都直接读取自反编译函数体。

GetPaddedRowCount — granule clamp

GetPaddedRowCount（0x13c90280）计算 SegmentedScan 操作的 per-segment 固定 stride：它把 minibatch row count 向上取到 int32 word 的 granule，并以 per-target minimum 为下限。反编译函数体短到可以精确呈现：

// GetPaddedRowCount (0x13c90280) — decompile-exact
int GetPaddedRowCount(const Target &t, int num) {
    uint64_t granule = t.GranuleBytes();                 // vtable[+0x5c0] (0x1d617f80)
    if (!t.SupportsSparseCore())                         // vtable[+0x260]
        LOG(FATAL) << "SparseCore is not supported by this target";  // target.h:1709
    uint64_t r = granule >> 2;                           // GranuleBytes / sizeof(int32) = i32 words/granule
    if (num > (int)r) r = num;                           // r = max(granule/4, num)
    int cfg = *(int*)( *((long*)&t + 297) + 148 );       // t[+0x948] -> [+0x94]  (per-target min-row config)
    if ((int)r <= cfg) r = cfg;                          // r = max(r, cfg)   <-- floor, not cap
    return (int)r;
}
```text

因此 padded row count 是 `max( max(GranuleBytes/4, num), cfg )`，即 per-minibatch row count 上的 granule-aligned floor。

> **注意 — config bound 是 floor（`max`），不是 cap（`min`）。** 反编译比较是 `if ((int)r <= cfg) r = cfg`，它在 `r` 低于 `cfg` 时将 `r` 向*上*提升到 `cfg`，因此是**下**界。padded count 因而是 `max(max(GranuleBytes/4, num), cfg)`。`cfg` 字段（`[Target+0x948]→[+0x94]`）的语义身份是从结构读取的：它是 per-target row-count config，但其 proto 字段*名称*未解码。

### 前向 / 反向 pass 名册

对于每个匹配的 op，`DecomposeForwardPass`（`0x1348a9c0`）及其梯度孪生 `DecomposeBackwardPass`（`0x1348b600`）发出 per-minibatch 内部 op。每个 minibatch 的前向 pass 会：通过 `GetNumSparseCores`（`0x13c9eba0`）+ `GetCores`（`0x14b79900`）读取 core 数量；从 `SparseDenseMatmulConfig` globals（`0x223a95b8`，字段 `+0x30`）读取 division-level 字段；调用 `CreateDynamicSliceCsr`；通过 `CreateCustomCall` 发出内部 `SparseDenseMatmulOp`；并用 `CreateGetTupleElement` 取得 per-minibatch 结果。顶层 driver 分派两种发出模式：

| Pass | 地址 | 作用 |
|---|---|---|
| `RunImpl` | `0x1348f940` | 扫描 op，构造 ArgSpecs，分派 |
| `IsSparseDenseMatmulWithMinibatchingOp` | `0x13c86da0` | `IsCustomCall(name)` 识别器 |
| `GetArgSpec` | `0x13c87040` | → `ForwardPassArgSpec` / `BackwardPassArgSpec` |
| `GetSparseDenseMatmulOpCustomCallTarget` | `0x13c86e60` | minibatching → 分解后的内部 op 名称 |
| `CreateDynamicSliceCsr` | `0x13489ea0` | per-minibatch `{base, padded}` 元组 |
| `GetPaddedRowCount` | `0x13c90280` | granule clamp |
| `CreateAddComputation` | `0x13c90320` | 构造 `"add"` reduction HLO computation |
| `DecomposeForwardPass` | `0x1348a9c0` | per-minibatch 前向发出 |
| `DecomposeBackwardPass` | `0x1348b600` | per-minibatch grad 发出 |
| `DecomposeForwardPassesWithNoWhileLoop` | `0x1348c720` | **inline** 前向（minibatches 展开） |
| `DecomposeForwardPassesWithWhileFusion` | `0x1348cd60` | **while-loop** 前向 |
| `DecomposeBackwardPassesWithNoWhileLoop` / `…WhileFusion` | `0x1348ca40` / `0x1348e260` | 反向对应形式 |
| `DecomposeSparseDenseMatmulWithMinibatchingWithWhileFusion` | `0x1348f200` | 顶层 while-fusion driver |
| `CombineParamsIntoTupleAndUpdateOutputShape` | `0x13c87260` | while-carry tuple（`CreateTuple`） |
| `CreateInitialWhileLoopInductionVar` | `0x1348a760` | while induction var |

> **陷阱 — 两种发出模式，同一描述符。** **no-while** 模式（`0x1348c720`）会 inline 发出每个 minibatch 的内部 op（展开）；**while-fusion** 模式（`0x1348cd60`）将它们包装在 HLO `while` loop 中，其 carry tuple 由 `CombineParamsIntoTupleAndUpdateOutputShape`（`0x13c87260`）组装，induction var 由 `CreateInitialWhileLoopInductionVar`（`0x1348a760`）播种。两者都消费相同的 `CreateDynamicSliceCsr` `{base, padded}` 描述符；模式只改变 per-minibatch 函数体在图中的布局，不改变 per-minibatch slice 算术。重新实现者必须产生两种形式，否则就要接受固定 trip count `while` 是大 batch 的 canonical 形状。

---

## 单元 2 — Operand Partition

### per-core CSR base offset

`CreateDynamicSliceCsr` 中乘/加链的全部目的，是为每个物理 SparseCore 提供进入单个拼接 `concatenated_csr_pointers` operand 的**连续、padded 窗口**。窗口起点是 `base` 值，长度是 `padded`。已逐字节确认的寄存器身份数据流如下：

| 符号 | 表达式 | 来源 |
|---|---|---|
| `padded` | `max( max(GranuleBytes/4, num), cfg )` | `GetPaddedRowCount`；`num` = arg int |
| `C_pad` | const `s32` = `padded` | `CreateConstant` |
| `GCI` | custom-call `"GetCoreIndex"` → `s32[1]` | runtime per-physical-SC index（op-type `0xc`） |
| `mul1` | `b * C_pad` | `b` = `HloInstruction*` arg（≈ minibatch index，**推断**） |
| `mul2` | `GCI * mul1` = `GCI · b · padded` | |
| `mul3` | `C_pad * GCI` = `padded · GCI` | |
| `base` | `mul2 + mul3` = `GCI·b·padded + padded·GCI` = `padded · GCI · (b + 1)` | per-core CSR 窗口 base offset |

因此每个物理 SparseCore 会读取拼接 row-pointers 中从 `padded · GCI · (b + 1)` 开始、长度为 `padded` 的连续窗口。`csr` operand（一个 `HloInstruction*` arg）存在时作为 `"DynamicSliceCsr"` custom-call 三个 operand 中的第一个被消费，`{sliced-csr, base, padded}` 被包装到返回的三元组中。

> **注意 — `b` 的 operand 身份是结构性的。** `CreateDynamicSliceCsr` 的两个 `HloInstruction*` 参数（称为 `a` 和 `b`）从 `DecomposeForwardPass` 调用点寄存器约定读取（`r8`=`a`，`r9`=`b`），而不是从具名 accessor 读取。`b` 是 minibatch index operand，`a` 是 csr operand。`CreateDynamicSliceCsr` 内的寄存器身份数据流逐字节精确；哪个 SSA 值是 per-table 还是 per-minibatch index，是从结构读取的。前向 pass 读取的 `SparseDenseMatmulConfig` 字段 `[+0x30]`（≈ `FLAGS_xla_sparse_core_minibatch_max_division_level`，`0x222bd280`）是一个 byte read，其 proto 字段*名称*来自 flag 邻接关系。

### 分解后的 operand-name 表

分解产生的 operand 布局由两个 `.data.rel.ro` 名称表固定，并已解析 relocation。前向 pass 有 7 个 operand，反向 pass 有 9 个（entry 8 重复 csr pointers）。Operand 0，`concatenated_csr_pointers`，是 `CreateDynamicSliceCsr` 按 minibatch 切片、并最终由 SegmentedScan 归约的 segment-id 来源：

```text
ForwardPassArgSpec::kForwardPassOperandNames (0x21937d80, 7 entries):
  0  concatenated_csr_pointers          <- segment-id source (per-minibatch sliced by CreateDynamicSliceCsr)
  1  concatenated_embedding_ids         (gather indices = sorted_token_ids)
  2  concatenated_sample_ids            (output rows)
  3  concatenated_gains                 (combiner weights — the per-id gain the DotCombiner applies)
  4  num_mini_batches_per_sparse_core   (scalar)
  5  embedding_table
  6  activations_init

BackwardPassArgSpec::kBackwardPassOperandNames (0x21938320, 9 entries):
  0..4  same as forward
  5     tables
  6     gradients
  7     hyperparameters                 (SGD / Adam / Ftrl / Adagrad / AdagradMomentum families)
  8     concatenated_csr_pointers        (repeated)

四个 concatenated_* operand 直接映射到 SC embedding sum-lookup 角色：csr_pointers → segment boundaries（SegmentedScan 的 reset points），embedding_ids → gather indices，sample_ids → output rows，gains → DotCombiner FMA 原样应用的 per-id 乘法权重。

表	地址	条目
kForwardPassOperandNames	0x21937d80	7
kBackwardPassOperandNames	0x21938320	9

---

单元 3 — Packed-Operands Lowering（SegmentedScanOp 重建器）

内部 SparseDenseMatmulOp op 存在后，它们降低到的 SC dialect sparse_core::SegmentedScanOp 必须在 packed-width operands 上重新创建（SC 向量引擎会 pack sub-byte / bf16 数据）。PackedOperandsLowering（runOnOperation 0x135d8520，ctor CreatePackedOperandsLoweringPass 0x135d82a0）是执行此操作的 MLIR full-conversion pass。它用 per-op dynamic-legality callback（setLegalityCallback 0x1c957e40 / 0x1c958640）和 RewritePatternSet 构造 ConversionTarget，然后运行 mlir::applyFullConversion（0x1c958ac0）。一个 op 合法当且仅当其 operands 已经处于目标 packed width；否则匹配到的 rewrite 会将其包装为 Unpack → op → Pack。

SegmentedScanOp 分支由 AddDynamicallyLegalScanOps<SegmentedScanOp>（legality lambda 0x135f3920）→ ScanOpLowering<SegmentedScanOp,SegmentedScanOp>::matchAndRewrite（0x135f3000）注册。逐字节确认的 rewrite 函数体如下：

ScanOpLowering<SegmentedScanOp>::matchAndRewrite (0x135f3000)  — decompile-confirmed
  1   ReductionOp = SegmentedScanOp::getReductionOp()           // 0x145fd460 — read reduction_op StringAttr
  2   UnpackOperand<UnpackFOp>(data, …)                         // 0x1360fac0 — split bf16/sub-byte data (F path)
  2'  UnpackOperand<UnpackUIOp>(segment-id, …)                  // 0x136104e0 — split (unsigned-int path)
  3   newop = SegmentedScanOp::create(b, loc, T, data, seg, reduction_op)  // 0x145fd5a0 — re-create on packed ops
  4   PackResults<PackFOp>(…)                                   // 0x13610940 — re-pack F results
  4'  PackResults<PackUIOp>(…)                                  // 0x13610de0 — re-pack UI results
```text

这是 dialect `SegmentedScanOp` **builder**：它产生 SC scan lowering（segmented-scan emission，[scan datapath](scan-datapath.md)）随后转换为 sequencer intrinsic 的 op。同一个 pass 会 legalize 约 40 个其他 AluEp arith/math op（`AddF/I`、`MulF/I`、`DivF`、`Max/Min`、`CmpF/I`、`Exp`、`Tanh`、`Rsqrt`、`Clamp` 等），每个都有自己的 `Unpack{F,SI,UI}` / `Pack{F,SI,UI}` 对；`SegmentedScan`/`Scan` 是该 packing 表中的两个成员。普通（非 segmented）`ScanOp` 分支是 `0x135f2580`（legality lambda `0x135f2f60`）。

> **注意 — partition 喂给 scan，而不是反过来。** CSR partition（单元 2）决定每个 minibatch 的 SegmentedScan 在*哪些* row-pointers 上归约；`PackedOperandsLowering` 决定该 SegmentedScan 的 operands 具有*什么 packed width*。二者正交：partition 是 HLO 层 index 算术，packing 是 dialect 层 operand-width legalization。重新实现者先运行 minibatching decomposition（产生基于 CSR slices 的 per-minibatch 内部 op），再运行 packed-operands lowering（在 packed operands 上重新创建每个 SegmentedScan）。

### activations / gradients layout adapter

`EmbeddingDataFormattingDecomposer`（`RunImpl` `0x1368b4a0`）是第三个 pass，也是最容易与 operand partition 混淆的一个。它**不是** CSR/id/gain 重排器。它在 dense XLA tensor layout 与 SC packed stacked-table layout 之间重排 dense per-table *activations*（前向输出）和 *gradients*（反向输入）。它匹配四个 `SparseCoreOperationType` custom-call：

| op-type | name | 方向 |
|---|---|---|
| `0x1a` | `SparseActivationsUnstack` | SC packed → per-table dense（前向输出） |
| `0x1b` | `SparseActivationsUnstackInterleaved` | 同上，interleaved |
| `0x1c` | `SparseGradientsStack` | per-table dense → SC packed（反向输入） |
| `0x1d` | `SparseGradientsStackInterleaved` | 同上，interleaved |

每个分派都基于 `EnableEmbeddingDataFormattingOffload(GetTpuCompEnv(op))`（`0x1d6b94a0` / `0x1d73de80`）：`true` → **Sc**（on-device SparseCore）变体；`false` → **Tc**（host/TensorCore）变体。Sc unstack（`DecomposeActivationsUnstackSc` `0x13682d40`）通过 `CreateSlice` + `CreateReshape` + `CreateConvert` + `CreateTuple` 将 packed stacked-table activation 拆成 per-table dense rows，由 `StackedTableConfig::Extract`（`0x13681160`）驱动，并按 `ElementPackingFactor`（`0x1d6b03e0`）× `NumEmbeddingDevices`（`0x1d6b8a00`）定尺寸。Sc stack（`DecomposeGradientsStackSc` `0x13684d40`）是逆过程：per-table dense grads → packed stacked grad，通过 `CreateConcatenate` + `CreatePad`（pad 到 stacked extent）+ `CreateSlice` + `CreateConvert` + `CreateUnary`。

> **陷阱 — 此 pass 不喂给 SegmentedScan operands。** 很自然会以为 "data formatting decomposer" 会重排 SegmentedScan 读取的 CSR/id/sample/gain operands。它不会。CSR→segment-id 来源是 `MinibatchingDecomposition`（单元 1）。`EmbeddingDataFormattingDecomposer` 只触碰 activations/gradients（dense ↔ packed-stacked-table layout adapter）。要把两个 pass 的 operand 域分开。

---

## 单元 4 — 绑定：从内部 op 到 Stream-op DAG

分解发出的每个 per-minibatch 内部 `SparseDenseMatmulOp` 本身会被降低为 SC-dialect **gather → sort → uniquify → segmented-reduce → scatter** Stream-op DAG。该 lowering *不是*由上述三个 pass 中任何一个完成，而是 `SparseDenseMatmulDotCombinerEmitter::Emit`（`0x1332bda0`，通过 `LoweringEmitter::EmitSparseDenseMatmulDotCombiner` `0x131a7ca0`）的工作；它分为 `EmitValencyLoop`（`0x1332cee0`）、`EmitVectorizedLoop`（`0x1332e1c0`）和 `EmitSampleCombiner`（`0x1332c640`）。它产生的 DAG：

| stage | SC dialect op | 作用 |
|---|---|---|
| gather | `IndirectStreamStartOp` / `IndirectVectorStreamStartOp`（build `0x145cf440`） | indirect HBM→SPMEM embedding-row load keyed by sorted token ids |
| sort | `SortOp`（build `0x14604480` / `0x146046c0`） | `(sample, token)` ids 的 lexicographic sort |
| uniquify | `UniqueOp` / `UniqueWithLaneIdsOp` + `DuplicateCountOp` / `…WithLaneIdsOp` | 去重 token ids，按 lane 计数重复项 |
| reduce | `SegmentedScanOp`（`reduction_op="sum"`） | per-segment sum-scan，在每个 CSR 边界 reset |
| scatter | `IndirectStreamAddStartOp` / `IndirectVectorStreamAddStartOp`（build `0x145d4420`） | scatter-ADD 到 HBM（前向 drain / 反向 grad accumulate） |

链条是：**`MinibatchingDecomposition`** 产生基于 CSR slices 的 per-minibatch 内部 op → **`PackedOperandsLowering`** 在 packed operands 上重新创建每个 `SegmentedScanOp` → **`DotCombiner` emitter** 在该 scan 周围连接 gather/sort/uniquify/scatter → `LowerToSparseCoreLlvm` 将 dialect 降低为 sequencer program。`concatenated_gains` operand（单元 2 operand 3）是 [DotCombiner FMA](sample-combiner-emitter.md) 应用的 per-id 乘法权重；`concatenated_csr_pointers`（operand 0）是 [valency loop](emit-valency-loop.md) 作为 per-sample trip count 读取的 segment boundary。

另有一个用于 gather-mul-scatter 形式的非 minibatch 路径：`GatherMulScatterSparseDenseMatmulOpDecomposer`（`0x13c861e0`）及 `GatherEmitter::ScatterOperandSlicesToHbm{,ForSortedIndices,ForChunkGather,ForColumnWiseGather}` 家族（`0x138e6c80` …）。那是对称的非 CSR lowering，不在本文范围内。

> **注意 — 深层 emitter wiring 在这里是概览，不做指令级解码。** `EmitVectorizedLoop` / `EmitSampleCombiner` 内部的逐指令 gather/scan/scatter wiring，包括每种 dtype 使用哪个 `IndirectStream` 变体、SPMEM tile 大小、`Sfence`/`TileBarrier`/`SyncAdd` 布置，以及 `GetScheduleType`（`0x131d5300`）的 valency-vs-vectorized 选择，由 [EmitValencyLoop](emit-valency-loop.md) 和 [SampleCombiner Emitter](sample-combiner-emitter.md) 负责。DAG *形状*（上述五个 stage）由 op `::build`/`::create` 的存在建立；loop 内部指令序列记录在那些页面上。

---

## 交叉引用

- [SparseCore Overview](overview.md) — Part IX 的导航入口；引擎名称、各代存在性，以及此分解打开的 embedding data path。
- [SparseCore Hardware Architecture](architecture.md) — 引擎职责、4:1 SC:TC 比例，以及 `GetCoreIndex` 索引的 physical-core geometry。
- [SC Backend Pipeline](sc-backend-pipeline.md) — 本文产生的 dialect op 会经过的十二 pass SC-MLO pipeline（以及 MEGACORE barrier）。
- [SC Core Selection](sc-core-selection.md) — 决定每个 collective 占用哪些 SC core 的 physical-core selection policy（对应本文的 per-core `GetCoreIndex` partition）。
- [SampleCombiner Emitter](sample-combiner-emitter.md) — 降低每个内部 `SparseDenseMatmulOp` 的 per-sample gather-multiply-accumulate emitter；消费本文 partition 的 `concatenated_gains` operand。
- [EmitValencyLoop](emit-valency-loop.md) — per-id scalar loop，其 trip count 是 `concatenated_csr_pointers` operand 定义的 per-sample CSR segment length。
- [SC Scan Datapath](scan-datapath.md) — 重建后的 `SegmentedScanOp` 后续降低到的 segmented-scan emission。
- [Stream Gather / Scatter](stream-gather-scatter.md) — 内部 op lowering 发出的 `IndirectStream*StartOp` gather/scatter-add primitive。
- [getSequencerType](getsequencertype.md) — lowered bundle 路由经过的 SCS/TAC/TEC engine-selection 函数。
- **Binary:** `extracted/libtpu-0.0.40-cp314-cp314-manylinux_2_31_x86_64/libtpu/libtpu.so`（build-id `89edbbe81c5b328a958fe628a9f2207d`）
- **索引条目:** Part IX — SparseCore & BarnaCore / SparseCore datapath (embeddings) — [返回索引](../index.md)