分段扫描
本页中的每个地址、reduction-string XOR 常量、vtable slot、struct offset、operand name 和 error string，都按字节精确读取自
libtpu-0.0.40-cp314wheel 中的libtpu.so（build-id89edbbe81c5b328a958fe628a9f2207d；构建libtpu_lts_20260413_b_RC00，未 strip）— 来源包括SegmentedScanOpLowering::matchAndRewrite函数体、SegmentedScanOp::build/create/getReductionOp、FindAndEmitToUnusedPort<…SegmentedAddScanF32>分配器、XlaSparseDenseMatmulWithCsrInputOp::Compile和MinibatchingDecomposition::CreateDynamicSliceCsr。.text/.rodataVMA == file offset（base0xe63c000/0x84a0000）；.data.rel.roVMA−0x200000== offset（reloc addends 通过readelf -rW读取）。地址适用于此构建；其他版本会不同。
摘要
SegmentedScanOp 是一种 prefix scan，它会在每个 segment 边界重置运行中的累加器，因此对 packed ragged batch 做一次硬件扫描即可为每个 segment 生成独立 prefix。它是 SparseCore embedding-sum lookup 的 reduce 原语：从密集 embedding 表 gather 出来的行被拼接成一个长向量，而 segment-boundary 操作数 — sparse minibatch 的 CSR row-offsets — 把该向量划分成逐样本 run。每个 run 上的 inclusive scan，在其最后一个 lane 取值，就是求和后的 embedding 行。本页记录 MLIR SegmentedScanOp lowering（0x13589d40）、驱动重置的 segment-boundary operand 绑定，以及供给它的 XlaSparseDenseMatmulWithCsrInput HLO custom-call 链。
该 lowering 几乎是 Scan Datapath 中 plain ScanOpLowering 的孪生版本，而且这是有意为之：它复用了完全相同的 3 字符 reduction-string XOR switch（sum/min/max）和完全相同的 element-type 轴（{i32, f32, i16, bf16}，通过 identity 与 builder 的 canonical types 比较）。结构差异有两点。第一，它绑定第二个 SSA operand — segment-id vector — 而 plain scan 在此处绑定 per-lane vector mask；SegmentedScanOp::build（0x145fd4a0）无条件发出两次 addOperands 调用，顺序为 (data, segment)。第二，它没有 i1/mprefix count-active 路径（segment scan 归约数据值；它不对 predicate bit 做 population count）。发射的 intrinsic 是 segmented 家族 tpu_*_seg_scan* — 七个不同的 codegen leaf，而不是 plain-scan 的九个 leaf。
决定性的重实现事实是：segment boundary 是一个普通 V read-port operand，不是 side register，也不是固定 port。在 ISA 发射时，data operand 和 segment-id operand 分别由 7 项 greedy first-free 分配器（FindAndEmitToUnusedPort，0x13ab2aa0）分配最低空闲 port；segment-id 会落在周围 bundle port pressure 下可用的任意 port。相比之下，per-lane vector mask 是一个独立 bundle 字段（proto+0x38），与 plain scan 完全相同。重实现者如果把 segment-id 固定到某个 V-index，或把它与 mask 混同，就会错误建模 operand frame。本页追踪完整 HLO → dialect → intrinsic → ISA 链，以便重实现者可以从 CSR custom-call 一路重建到 per-segment-reset scan 的 embedding sum-lookup。
对于重实现，契约如下：
- lowering 逐字复用 plain-scan reduction switch，然后绑定第二个 operand。 reduction 通过常量 XOR 从 3 字符
StringRef解码（sum=0x7573|0x6d，min=0x696d|0x6e，max=0x616d|0x78），element type 通过 identity-compare 与getI32Type/getF32Type/getI16Type/getBF16Type比较。没有 enum，没有strcmp。与 plain scan 相同；差异是 operand frame 和发射的 intrinsic family。 operand[0]= data，operand[1]= segment boundary。SegmentedScanOp::build无条件添加二者（没有if(data)guard，不同于ScanOp::build）。segment vector 是 reset signal：segment id 改变处，累加器重新开始。- 发射七个 intrinsic，全部为
NOperands<2>。tpu_add_seg_scan1xN{i,f}、tpu_add_half_seg_scan2xN（i16/bf16共享这个 packed-pair arm）、tpu_min_seg_scan1xN{i,f}、tpu_max_seg_scan1xN{i,f}。tpu_{min,max}_seg_scan2xN已注册但从不 codegen — 没有::create。不存在i1/mprefixsegmented 路径。 i16/bf16segmented-add 由 target capability 门控。sum×{i16,bf16}arm 调用 target subobject 上的 vtable slot+0x780（(**(ctx+0x68) + 1920)(…)），若为 false，则发出"Currently seg scan add for bf16 is only supported"并失败。没有 bf16 ALU 的代际上，重实现者必须拒绝 half-precision segmented sum。- segment boundary 占用一个空闲 V read port，而不是固定 index 或 mask。 7-port greedy first-free 分配器（slots
+0x1c..+0x34，present mask+0x10）为它分配任何打开的 port。per-lane mask 是独立的proto+0x38字段。operand routing 必须按 op identity 分支。 - CSR row-offsets 成为 segment-id。
XlaSparseDenseMatmulWithCsrInputOp::Compile发射SparseDenseMatmulWithMinibatchingOpcustom-call；MinibatchingDecomposition将concatenated_csr_pointers切成逐 minibatch offset vectors；这些 offsets 就是SegmentedScanOp的operand[1]。
| MLIR op | mlir::sparse_core::SegmentedScanOp（SC dialect；reduction_op 3 字符 StringAttr） |
| Lowering | SegmentedScanOpLowering::matchAndRewrite 0x13589d40 — reduction × dtype → tpu_*_seg_scan* |
| Build / create | SegmentedScanOp::build 0x145fd4a0（先 addOperands(data)，再 addOperands(segment)）；::create 0x145fd5a0 |
| Reduction 读取 | SegmentedScanOp::getReductionOp 0x145fd460 — property word ((w>>19)&0x10)+64 → StringAttr::getValue |
| Result drain | LLVMStructType::getLiteral 0x17471ae0 {value, segment-id} → LLVM::ExtractValueOp 0x1728c5a0 |
| BF16 gate | target subobject (ctx+0x68) 上的 vtable +0x780（1920）；错误 "Currently seg scan add for bf16 is only supported"（.rodata 0x87036bf，49 B） |
| ISA emit | EmitVectorResultUnop<…SegmentedAddScanF32>（gfc 0x13aaf560）；port alloc FindAndEmitToUnusedPort 0x13ab2aa0（gfc）/ 0x13a4b680（glc） |
| ISA op (f32) | SparseCoreTecVectorExtended_SegmentedAddScanF32，proto inst oneof 0x23 |
| Dialect rewrite | PackedOperandsLowering ScanOpLowering<SegmentedScanOp> 0x135f3000（unpack → re-create → pack） |
| HLO front-end | XlaSparseDenseMatmulWithCsrInputOp::Compile 0xe650800 → custom-call SparseDenseMatmulWithMinibatchingOp（35 B），7 个 operands |
| CSR → segment-id | MinibatchingDecomposition::CreateDynamicSliceCsr 0x13489ea0 切片 concatenated_csr_pointers |
| Confidence | 已确认（反编译锚定），除非某行或 callout 另有说明 |
注意 — 本页负责
SegmentedScanOplowering、segment-boundary reset 和 CSR matmul 链。 plain-scan mask datapath、两个 M-register band、post-scanVectorSelect以及i1/mprefixcount path 位于 Scan Datapath，此处不重复。SegmentedAddScanISA operand frame、VpackFormatcapability matrix 和完整 segmented-scan proto oneof case map 位于 Segmented Add-Scan。VEX bundle bit positions 位于 VEX Mask/Dest-Port/Sub-Opcode。它们只交叉链接，不重复。
Segment-Boundary Reset 模型
目的
先固定语义模型，再进入 lowering 细节，因为它正是 segmented scan 与 Scan Datapath 中 plain scan 的区别。对 lane vector x[0..N) 的 plain inclusive scan 生成 y[i] = x[0] ⊕ x[1] ⊕ … ⊕ x[i]，其中 ⊕ 是结合性 reduction。Segmented scan 增加一个并行的 segment-id vector s[0..N)：当相邻 lane 之间的 segment id 改变时，累加器重置为 reduction identity，因此 scan 永远不会把值跨 boundary 携带。
重置规则
text
Segmented inclusive scan — the accumulator restarts at each boundary
data x: [ a0 a1 a2 | b0 b1 | c0 c1 c2 c3 ]
segment s: [ 0 0 0 | 1 1 | 2 2 2 2 ] (CSR-derived segment ids)
^boundary ^boundary
inclusive: [ a0 a0⊕a1 a0⊕a1⊕a2 | b0 b0⊕b1 | c0 c0⊕c1 … c0⊕c1⊕c2⊕c3 ]
└── per-segment prefix; the carry does NOT cross a '|' ───────┘
reset rule: acc[i] = (s[i] == s[i-1]) ? acc[i-1] ⊕ x[i] // continue the run
: identity ⊕ x[i] // s changed → restart
```text
对 embedding sum-lookup 而言，reduction `⊕` 是 `add`，每个 segment 是一个样本的 gathered embedding rows bag。某个样本的 summed embedding 是其 segment **最后一个 lane** 上的 inclusive scan 值 — 即 per-segment total。CSR row-offsets 精确定义一个样本 run 在何处结束、下一个样本 run 在何处开始，因此 offsets *就是* segment ids（见 [CSR Matmul 链](#csr-matmul-链)）。
> **怪癖 — segment operand 重置累加器；它不 mask lanes。** 读过 [Scan Datapath](scan-datapath.md) 的读者知道 plain scan 的 `operand[1]` 是 per-lane M-register mask，用来门控*哪些 lanes 参与*。segmented scan 的 `operand[1]` 是一个*值*向量，硬件逐 lane 比较它来决定*在哪里重新开始 carry*。它们占据相同 SSA slot index，但含义相反，并路由到不同 bundle 字段。segmented scan 上仍然存在 mask（独立的 `proto+0x38` 字段），但它不是 segment operand。
### Inclusive 与 exclusive
发射的 intrinsic 是 inclusive form（`tpu_*_seg_scan*`）：位置 `i` 包含 `x[i]`。exclusive 变体 — `y[i] = identity ⊕ x[0] ⊕ … ⊕ x[i-1]`，排除 `x[i]`，在每个 segment 的第一个 lane 重置为 identity — 在此构建中不是独立 intrinsic；exclusive segmented scan 由 front-end 通过移动 inclusive result 合成（exclusive seed 是 boundary 处的 per-segment identity）。硬件原语是 inclusive。（shift-to-exclusive rewrite 是 front-end pattern，不在本页；其精确 lowering 为 LOW。）
---
## MLIR Lowering
### 目的
`SegmentedScanOpLowering::matchAndRewrite`（`0x13589d40`）将 `sparse_core::SegmentedScanOp` 重写为一个 segmented-scan intrinsic，由 reduction string 和 result element type 选择，然后 drain `{value, segment-id}` result struct。它是一个 `ConvertOpToLLVMPattern`；dispatch 函数体是扁平的 reduction-string-XOR → element-type cascade，结构上与 plain `ScanOpLowering`（`0x1358ab00`）相同，只是去掉 `i1` arm，并增加 segment operand。
### 入口点
```text
sparse_core::SegmentedScanOp (NOperands<2>, reduction_op StringAttr)
└─ SegmentedScanOpLowering::matchAndRewrite (0x13589d40) ── reduction × dtype → intrinsic
├─ SegmentedScanOp::getReductionOp (0x145fd460) ── 3-char StringRef
├─ VectorType::getElementType ── result element type
├─ Builder::get{I32,F32,I16,BF16}Type (0x1d853c40 / 0x1d853980 / 0x1d853c20 / 0x1d853680)
├─ VectorType::get(i1, …) (0x1d894100) ── per-lane boundary mask type
├─ LLVMStructType::getLiteral (0x17471ae0) ── {value, segment-id} result pair
├─ tpu_*_seg_scan*::create ── the 7 emitted leaves (TABLE)
└─ LLVM::ExtractValueOp::create (0x1728c5a0) ── pull value(idx0)/segment-id back算法
reduction string 通过对 3 个字节做常量 XOR 来解码 — len == 3 然后 (word0 ^ K0) | (byte2 ^ K1) == 0。element type 从转换后的 result VectorType 中取得，并通过 pointer-identity 与 builder 的 canonical types 比较。这里没有 i1 特例 — 第一个分支是 sum，不是 count-active 路径：
c
function SegmentedScanOpLowering_matchAndRewrite(op, adaptor): // 0x13589d40
red = op.getReductionOp() // StringRef, 3 chars (0x145fd460)
elt = getElementType(converted_result_type) // i32 / f32 / i16 / bf16
// --- sum: (word0 ^ 0x7573) | (byte2 ^ 0x6d) == 0 --- (line 84)
if len(red) == 3 && red == "sum":
if elt == i32: emit tpu_add_seg_scan1xNi // 0x146d5c40 (line 254)
elif elt == f32: emit tpu_add_seg_scan1xNf // 0x146d5a80 (line 273)
elif elt in {i16, bf16}:
// BF16/I16 capability gate — vtable +0x780 (1920) on target subobject (ctx+0x68)
if !(**(ctx+0x68) + 1920)(ctx+0x68): // line 289
emitError("Currently seg scan add for bf16 is only supported") // 0x87036bf, 49 B
return failure
emit tpu_add_half_seg_scan2xN // 0x146d45c0 (line 302)
else: return failure
// --- max: (word0 ^ 0x616d) | (byte2 ^ 0x78) == 0 --- (line 87)
elif len(red) == 3 && red == "max":
if elt == f32: emit tpu_max_seg_scan1xNf // 0x14730e00 (line 118)
elif elt == i32: emit tpu_max_seg_scan1xNi // 0x14730fc0 (line 137)
else: return failure // no i16/bf16 max-seg arm
// --- min: (word0 ^ 0x696d) | (byte2 ^ 0x6e) == 0 --- (line 142)
elif len(red) == 3 && red == "min":
if elt == f32: emit tpu_min_seg_scan1xNf // 0x147316c0 (line 177)
elif elt == i32: emit tpu_min_seg_scan1xNi // 0x14731880 (line 213)
else: return failure // no i16/bf16 min-seg arm
else:
return failure // xor ebx,ebx — unknown reduction string
// --- result drain: the intrinsic returns an LLVMStructType{value, segment-id} ---
struct_ty = LLVMStructType::getLiteral({value_vec, i1_seg_vec}) // 0x17471ae0
value = ExtractValueOp(struct_result, idx=0) // 0x1728c5a0 (lines 218 / 394)
seg_id = ExtractValueOp(struct_result, idx=1)
replaceOp(op, value)
return success
```text
XOR immediates 是直接从 `cmp` 行读取的小端 byte triples：`"sum"` = `s u`=`0x7573`，`m`=`0x6d`；`"max"` = `m a`=`0x616d`，`x`=`0x78`；`"min"` = `m i`=`0x696d`，`n`=`0x6e`。这些与 plain `ScanOpLowering` 常量逐字节相同 — 两个 lowering 共享 reduction vocabulary。
> **注意 — reduction string 是 `sum`/`min`/`max`，不是 `add`/`min`/`max`。** `0x13589d8f` 处的 `matchAndRewrite` XOR test 与 `0x7573|0x6d` = `"sum"` 比较，因此 embedding sum-lookup 的 canonical reduction-kind string 是 `"sum"`。（字节确认；plain `ScanOpLowering` 使用完全相同的 immediate。）
### reduction × element-type → intrinsic 映射
每个 arm 都以字节锚定到 lowering 函数体中的具体 `tpu_*_seg_scan*::create` 调用点：
| reduction | result elt | → intrinsic | `::create` @ | gate |
|---|---|---|---|---|
| `sum` | `i32` | `tpu_add_seg_scan1xNi` | `0x146d5c40` | — |
| `sum` | `f32` | `tpu_add_seg_scan1xNf` | `0x146d5a80` | — |
| `sum` | `i16` | `tpu_add_half_seg_scan2xN` | `0x146d45c0` | `+0x780` must be true |
| `sum` | `bf16` | `tpu_add_half_seg_scan2xN` | `0x146d45c0` | `+0x780` must be true |
| `max` | `f32` | `tpu_max_seg_scan1xNf` | `0x14730e00` | — |
| `max` | `i32` | `tpu_max_seg_scan1xNi` | `0x14730fc0` | — |
| `min` | `f32` | `tpu_min_seg_scan1xNf` | `0x147316c0` | — |
| `min` | `i32` | `tpu_min_seg_scan1xNi` | `0x14731880` | — |
| any | other elt | `emitError` → failure | — | — |
`i16` 和 `bf16` 共享单个 `tpu_add_half_seg_scan2xN` arm（packed-pair `PartialSum` widen）。`min` 和 `max` **没有** `i16`/`bf16` 发射 arm：`tpu_{min,max}_seg_scan2xN` 注册为带 `NOperands<2>` trait 的 dialect ops，但**没有 `::create`/`::build`** — 已声明但不 codegen。lowering 从不生成 2xN min/max segmented scan；只有 `add` 具有 half/2xN widen。
> **注意 — 只有 `add` 具有 half-precision segmented widen，且它受 target 门控。** `i16`/`bf16` segmented-sum arm 是*唯一*触碰 `2xN` packed form 的路径，并且只有当 target 的 vtable slot `+0x780`（bf16-ALU / EUP lane-width capability）返回 true 时才合法。`(ctx+0x68)` 处的 subobject pointer 由 `LowerToSparseCoreLlvmPass::lowerFunc`（`0x13568280`）设置为 codegen target 的 `+0x8` sub-object。在没有 native bf16 lane 的代际上，lowering 发出 `"Currently seg scan add for bf16 is only supported"`（`.rodata` `0x87036bf`，49 bytes）— 运行时第二个片段 `" for GXC"`（8 bytes）会追加到 diagnostic stream 中，因此用户可见消息读作 `…only supported for GXC` — 然后失败，所以 SC bf16 segmented sum 限于拥有 bf16 ALU 的代际。此 gate 是 [Scan Datapath](scan-datapath.md) 中 plain scan 相同 `+0x780` 检查的 segmented-scan 孪生版本。
### 为什么没有 `i1` count path
plain `ScanOpLowering` 对 `i1`（boolean）输入特殊处理为 `tpu_mprefix` population-count-prefix 原语（见 [Scan Datapath](scan-datapath.md#the-i1-count-active-path)）。segmented lowering **没有**这样的 arm：它的第一个分支是 `sum`，element-type 轴只有 `{i32, f32, i16, bf16}`。segmented scan 归约按 segment 划分的数据*值*；它不在 boundary 上计数 predicate bits。移植 plain-scan dispatch 的重实现者必须为 segmented variant 删除 `i1` arm。
---
## Operand Frame
### 目的
lowering 行为只是故事的一半；另一半是两个 SSA operands 在 build time 如何绑定，并在 ISA emit 时如何路由。segment boundary 是第二个 operand，其绑定 — 包括 MLIR `build` 顺序和 ISA port allocation — 是重实现者必须精确复现的内容，才能把 reset signal 喂给硬件。
### Build 顺序 — data 第一，segment 第二
`SegmentedScanOp::build`（`0x145fd4a0`）接收 `(OpBuilder, OperationState, Type result, Value data, Value segment, StringAttr reduction)`，并按顺序发出**两次无条件** `addOperands` 调用：
```c
// SegmentedScanOp::build(OpBuilder, OperationState&, Type, Value data, Value segment, StringAttr red)
// (0x145fd4a0)
addOperands(state, &data, 1); // operand[0] = data (line 16)
addOperands(state, &segment, 1); // operand[1] = segment boundary (line 17)
state.properties.reduction_op = red; // StringAttr property对比 ScanOp::build（0x145f92e0，记录于 Scan Datapath）：它会guard data operand（if (data) addOperands(data)），然后把 per-lane vector mask 作为 operand[1] 添加。segmented build 没有 guard，并把 segment vector 绑定为 operand[1]。getReductionOp（0x145fd460）从 op 的 inline-or-out-of-line property word 读取 reduction_op StringAttr：offset ((property_word >> 19) & 0x10) + 64，然后 StringAttr::getValue。
坑点 —
operand[1]对ScanOp表示 vector mask，对SegmentedScanOp表示 segment boundary。 二者都是 SSAoperand[1]，但 plain scan 把它路由到 in-scan M-register mask 字段（proto+0x38），segmented scan 把它路由到 V read port（一个值 operand）。重实现者 wiring operand frame 时必须按 op identity 分支，而不能假设operand[1]总是 mask。per-lane mask 在 segmented scan 上仍然作为独立的proto+0x38字段存在；segment operand 是额外存在的，而不是替代它。
ISA emit — segment-id 占用空闲 V read port
在 ISA emit 时，EmitVectorResultUnop<…SegmentedAddScanF32>（gfc 0x13aaf560）读取 MCInst operands 并路由它们。per-lane mask 来自 operand[1] → GetVectorMask → 字段 proto+0x38，present-bit proto+0x11 |= 1（与 plain scan 相同）。data value 来自 operand[2] → GetVregno → FindAndEmitToUnusedPort — segment-id 也是如此：二者各自被分配下一个空闲 port。
FindAndEmitToUnusedPort（gfc 0x13ab2aa0，glc 0x13a4b680）是一个覆盖 7 项 unused SparsecoreVregReadPort（ports 0..6）btree_set 的 greedy first-free 分配器。它弹出最低空闲 port，然后 switch (port) 把解析出的 Vregno 写入七个连续 struct slots 之一，并设置匹配的 present-bit：
c
// FindAndEmitToUnusedPort<SparsecoreVregReadPort, …SegmentedAddScanF32> (0x13ab2aa0 gfc)
port = btree_set_pop_lowest(unused_ports); // erase the lowest free port
switch (port): // line 29
case 0: inst[0x1c] = vregno; inst[0x10] |= 0x02; break;
case 1: inst[0x20] = vregno; inst[0x10] |= 0x04; break;
case 2: inst[0x24] = vregno; inst[0x10] |= 0x08; break;
case 3: inst[0x28] = vregno; inst[0x10] |= 0x10; break;
case 4: inst[0x2c] = vregno; inst[0x10] |= 0x20; break; // line 53
case 5: inst[0x30] = vregno; inst[0x10] |= 0x40; break;
case 6: inst[0x34] = vregno; inst[0x10] |= 0x80; break;
default: return MakeError("Unsupported Port Value: $0"); // line 68 (isa_emitter_utils.h:3114)
// (the pre-switch guard rejects an empty pool before the pop)
if (unused_ports.empty()): // out of ports — pre-switch RET_CHECK
RetCheckFailSlowPath("!port_is_free.empty()"); // line 98 (isa_emitter_utils.h:3025)
statusor.port = port; // returned port enum
```text
| port id | struct slot (Vregno) | present-bit (`inst[0x10] \|=`) | returned enum |
|---|---|---|---|
| 0 | `+0x1c` | `0x02` | 0 |
| 1 | `+0x20` | `0x04` | 1 |
| 2 | `+0x24` | `0x08` | 2 |
| 3 | `+0x28` | `0x10` | 3 |
| 4 | `+0x2c` | `0x20` | 4 |
| 5 | `+0x30` | `0x40` | 5 |
| 6 | `+0x34` | `0x80` | 6 |
> **注意 — 有 7 个 V read ports，segment-id 不是固定 V1。** 反编译显示一个覆盖**七个** ports（0..6）的 `switch`，写入 slots `+0x1c..+0x34`，从 `btree_set` 中 greedy first-free 分配。segment-id 会落在周围 bundle port pressure 下可用的任意 port — 它的物理 V-index 是 allocator state 的函数，而不是常量。Vmask 是 `+0x38` 处的*独立*字段，与 port pool 不同。（字节确认；gfc 与 glc 共享该分配器。）
`SegmentedAddScanF32` ISA op 是 proto `inst` oneof case `0x23`（`0x13aaf600` 处的 accessor `mutable_segmented_add_scan_f32` 比较 `proto+0x58 == 0x23`）。bf16 form 是 `SegmentedAddScanBf16PartialSumBf16`（oneof `0x2e`）— gfc/glc proto pool 只有 `Bf16PartialSumBf16` form，没有更宽的 `PartialSumF32` segmented variant。完整的 18 项 segmented-scan oneof case map 和最终 instruction word 中的 per-port bit packing 由 [Segmented Add-Scan](segmented-add-scan.md) / [VEX Mask/Dest-Port/Sub-Opcode](vex-mask-destport-subopcode.md) 负责；本页只锚定 operand-routing 机制。
---
## CSR Matmul 链
### 目的
segmented scan 并不孤立出现 — 它是 SparseCore embedding sum-lookup 的 reduce 阶段，而它读取的 segment-id 是 sparse minibatch 的 CSR（compressed-sparse-row）row-offset vector。本节追踪从 `XlaSparseDenseMatmulWithCsrInput` HLO op 到 `SegmentedScanOp` 的链条，让重实现者能看清 *segment boundaries 从哪里来*。CSR 是标准 sparse-embedding 格式：`row_pointers[k]` 是 flat `(token_id, sample_id, gain)` arrays 中样本 `k` 的 looked-up rows bag 开始处的 offset，因此 row-pointers *就是* per-sample segment boundaries。
### 阶段 1 — HLO custom-call
`XlaSparseDenseMatmulWithCsrInputOp::Compile`（`0xe650800`）是 DLRM embedding sum-lookup 的 XLA op-kernel。它按顺序读取五个 named inputs，验证 minibatch count，构建 `FrontendAttributes` string-map backend-config，并发射一个带 7 个 operands 的 `xla::CustomCall`：
```c
function XlaSparseDenseMatmulWithCsrInputOp_Compile(ctx): // 0xe650800
row_pointers = ctx.Input("row_pointers") // line 124 — CSR offsets = segment boundaries
sorted_sample_ids= ctx.Input("sorted_sample_ids") // line 126 — per-id output row
sorted_token_ids = ctx.Input("sorted_token_ids") // line 128 — per-id embedding-table row (gather idx)
sorted_gains = ctx.Input("sorted_gains") // line 130 — per-id combiner gain (weight)
embedding_table = ctx.Input("embedding_table") // line 132 — the dense matrix gathered from
num_minibatches = ctx.Input("num_minibatches_per_physical_sparse_core") // line 226, validated scalar
fe = FrontendAttributes() // line 280
fe["_xla_compute_type"] = "sparse" // lines 287-288
fe["_xla_quantization_low_value"] = attr.quant_low
fe["_xla_quantization_high_value"] = attr.quant_high
fe["_xla_quantization_num_buckets_value"] = attr.quant_num_buckets
fe["_xla_enable_full_hbm_sort"] = "false" // lines 851-852 (default)
CustomCall(builder, "SparseDenseMatmulWithMinibatchingOp", // line 907; 35 B, 7 operands
operands={row_pointers, sorted_token_ids, sorted_sample_ids,
sorted_gains, embedding_table, num_minibatches, activations_init},
…, backend_config=fe)custom-call target name "SparseDenseMatmulWithMinibatchingOp"（35 bytes）由一个 32-byte .rodata prefix 加一个 4-byte "ngOp" tail inline 组装（strcpy(buf+31, "ngOp")，line 883）。第 7 个 operand（activations_init，accumulator initializer）按结构从 operand span 中读取，而不是来自单独的 Input() 调用（推断 = activations_init；匹配 decomposed forward operand[6]）。
| # | named input | role |
|---|---|---|
| 0 | row_pointers | CSR row-offsets = per-sample segment boundaries |
| 1 | sorted_sample_ids | per-id output (minibatch) row index |
| 2 | sorted_token_ids | per-id embedding-table row index (the gather index) |
| 3 | sorted_gains | per-id scale / combiner gain (weight) |
| 4 | embedding_table | the dense embedding matrix being gathered from |
| (scalar) | num_minibatches_per_physical_sparse_core | bounds the minibatch count (validated scalar) |
阶段 2 — minibatching decomposition 切片 CSR offsets
SparseDenseMatmulWithMinibatchingOp custom-call 由两个 SC HLO passes 重写：MinibatchingDecomposition（AddPass 0x1306d5c0）和 EmbeddingDataFormattingDecomposer（AddPass 0x1095b6a0）。MinibatchingDecomposition::CreateDynamicSliceCsr（0x13489ea0）逐 minibatch 切片并 pad 拼接后的 CSR row-pointers — 它构建 DynamicSlice、Binary（op 75 = add-style index arithmetic）和 CustomCall HLO ops，并通过 sparse_dense_matmul_decomposer_util::GetPaddedRowCount 约束切片：
c
function CreateDynamicSliceCsr(comp, csr_tuple, …, minibatch, …): // 0x13489ea0
RET_CHECK(csr_tuple.size() > kCsrTupleRowPtrIndex) // line 285
padded = GetPaddedRowCount(util, target, minibatch) // line 83
…
idx = CreateBinary(comp, 75, …) // lines 130/146/162 — slice index math
sliced = DynamicSlice(comp, concatenated_csr_pointers, idx, padded)
return CreateCustomCall(comp, …, sliced, …) // lines 112/262
```text
decomposed **forward** op 的 `operand[0]` 是 `concatenated_csr_pointers`（`ForwardPassArgSpec::kForwardPassOperandNames`，`.data.rel.ro` `0x21937d80`，reloc-resolved）。这些 per-sample row offsets *就是* segment boundaries：每个 output activation row 聚合一个连续 run 的 gathered embedding rows，而 run boundaries 就是 CSR offsets。
| # | decomposed forward operand | role |
|---|---|---|
| 0 | `concatenated_csr_pointers` | **the segment-id source** (sliced by `CreateDynamicSliceCsr`) |
| 1 | `concatenated_embedding_ids` | gather indices |
| 2 | `concatenated_sample_ids` | output rows |
| 3 | `concatenated_gains` | combiner weights |
| 4 | `num_mini_batches_per_sparse_core` | scalar |
| 5 | `embedding_table` | the dense matrix |
| 6 | `activations_init` | accumulator init |
backward（grad）op（`kBackwardPassOperandNames` `0x21938320`）在相同 CSR/id/sample/gain operands 之上增加 `tables` / `gradients` / `hyperparameters`（optimizer state — SGD/Adam/Ftrl/Adagrad/AdagradMomentum families），并在同一 CSR segment structure 上累加 per-segment gradient。
### 阶段 3 — dialect SegmentedScanOp 和完整 datapath
`PackedOperandsLowering` 的 `ScanOpLowering<SegmentedScanOp>`（`0x135f3000`）是*构建* dialect `SegmentedScanOp` 的 rewrite（通过 `SegmentedScanOp::create`，`0x145fd5a0`）：它 unpack `bf16`/sub-byte operands，使用 packed `(data, segment-id)` operands 和 `reduction_op` `StringAttr` 重新创建 `SegmentedScanOp`，然后 pack results。最终的 `SegmentedScanOpLowering`（[MLIR Lowering](#mlir-lowering) 中的函数体）再把它 lower 到 intrinsic。
```text
The complete embedding sum-lookup HLO → SC dialect → intrinsic → ISA datapath
HLO op-kernel XlaSparseDenseMatmulWithCsrInputOp::Compile (0xe650800)
└─ CustomCall "SparseDenseMatmulWithMinibatchingOp" (7 operands, frontend-attrs)
HLO pass A MinibatchingDecomposition::CreateDynamicSliceCsr (0x13489ea0)
└─ slice/pad concatenated_csr_pointers → per-minibatch segment-id vector
HLO pass B EmbeddingDataFormattingDecomposer (0x1095b6a0) ── activations stack/unstack
SC dialect sparse_core::SegmentedScanOp(data, segment-id=CSR-offsets, reduction_op="sum")
build 0x145fd4a0 / create 0x145fd5a0
dialect rewrite PackedOperandsLowering ScanOpLowering<SegmentedScanOp> (0x135f3000)
└─ unpack bf16 → re-create SegmentedScanOp → pack results
dialect→intr SegmentedScanOpLowering::matchAndRewrite (0x13589d40) ── reduction × dtype switch
└─ tpu_add[_half]/_min/_max_seg_scan{1xNf,1xNi,2xN} (+0x780 gate for bf16)
ISA op SparseCoreTecVectorExtended_SegmentedAddScan{dtype} (oneof add_f32=0x23)
EmitVectorResultUnop gfc 0x13aaf560 ; FindAndEmitToUnusedPort 0x13ab2aa0
└─ per-segment inclusive prefix-sum that RESETS at each CSR-offset boundary
result drain LLVM::ExtractValueOp (0x1728c5a0) ── value(idx0) / segment-id(idx1)注意 — CSR offsets 驱动 reset；gather 驱动 data。 sum-lookup 的两半映射到两个
SegmentedScanOpoperands：embedding_tablegather（由sorted_token_ids索引，并由sorted_gains缩放）产生 data vector（operand[0]）；row_pointers产生 segment-id vector（operand[1]）。对 gathered rows 做 inclusive segmented scan，并在每个 segment 的最后一个 lane 读取，就是 per-sample summed embedding。CreateDynamicSliceCsr中的确切 slice arithmetic（来自num_minibatches和GetPaddedRowCount的 start/stride）被观察为DynamicSlice+Binary+CustomCall，但 index expressions 未完全解码（literal transform 为 LOW）。
函数映射
| Symbol | Address | Role |
|---|---|---|
SegmentedScanOpLowering::matchAndRewrite | 0x13589d40 | reduction × dtype → tpu_*_seg_scan*；bf16 +0x780 gate；无 i1 path |
SegmentedScanOp::build | 0x145fd4a0 | 先 addOperands(data)，再 addOperands(segment) — operand[1]=boundary |
SegmentedScanOp::create | 0x145fd5a0 | 构建 (data, segment, reductionStr) |
SegmentedScanOp::getReductionOp | 0x145fd460 | property word ((w>>19)&0x10)+64 → StringAttr::getValue |
ScanOpLowering<SegmentedScanOp> (PackedOperands) | 0x135f3000 | unpack bf16 → re-create SegmentedScanOp → pack results |
LowerToSparseCoreLlvmPass::lowerFunc | 0x13568280 | 设置 (pattern+0x68)=(Target+0x8)，即 bf16-ALU capability holder |
EmitVectorResultUnop<…SegmentedAddScanF32> | 0x13aaf560 | gfc emit；op[1]→mask +0x38，op[2]→FindAndEmitToUnusedPort |
FindAndEmitToUnusedPort<…SegmentedAddScanF32> | 0x13ab2aa0 (gfc) / 0x13a4b680 (glc) | 7-port greedy first-free；slots +0x1c..+0x34，present +0x10 |
mutable_segmented_add_scan_f32 | 0x13aaf600 | proto inst oneof 0x23 accessor（proto+0x58 == 0x23） |
XlaSparseDenseMatmulWithCsrInputOp::Compile | 0xe650800 | 5 个 named inputs → SparseDenseMatmulWithMinibatchingOp custom-call，7 operands |
XlaSparseDenseMatmulWithCsrInputOp ctor | 0xe650140 | 读取 table_name/input_size/num_sc_per_logical_device/quant attrs |
GetMaxIdsAndUniques | 0xe651fa0 | 委托给 GetMaxIdsAndUniquesExternal — gather/dedup window bounds |
MinibatchingDecomposition::CreateDynamicSliceCsr | 0x13489ea0 | 逐 minibatch 切片 concatenated_csr_pointers（GetPaddedRowCount） |
EmbeddingDataFormattingDecomposer AddPass | 0x1095b6a0 | activations stack/unstack decomposition |
tpu_add_seg_scan1xNi / 1xNf ::create | 0x146d5c40 / 0x146d5a80 | sum i32 / f32 segmented leaves |
tpu_add_half_seg_scan2xN::create | 0x146d45c0 | sum i16/bf16 packed-pair leaf（gated） |
tpu_min_seg_scan1xNf / 1xNi ::create | 0x147316c0 / 0x14731880 | min f32 / i32 segmented leaves |
tpu_max_seg_scan1xNf / 1xNi ::create | 0x14730e00 / 0x14730fc0 | max f32 / i32 segmented leaves |
tpu_{min,max}_seg_scan2xN | (registered) | NOperands<2> trait，无 ::create — declared-but-uncodegen'd |
注意事项
- lowering 是 plain scan 减去
i1，再加第二个 operand。 复用 Scan Datapath 中的sum/min/maxXOR switch 和{i32,f32,i16,bf16}轴，删除i1/mprefixarm，把operand[1]绑定为 segment vector 而不是 mask，并发射tpu_*_seg_scan*family。 operand[1]是 reset signal，不是 mask。 它是一个值向量，硬件逐 lane 比较它以在 boundaries 处重启 carry。per-lane mask 是独立的proto+0x38字段，在每个 segmented scan 上仍然存在。operand routing 必须按 op identity 分支。- segment-id 占用一个空闲 V read port（7-port greedy allocator），不是固定 index。 Slots
+0x1c..+0x34，present mask+0x10。data 和 segment-id 各自占用下一个空闲 port；物理 index 取决于 bundle port pressure。 i16/bf16segmented-add 是唯一 half-precision arm，并且受 target 门控。 检查 target subobject 上的 vtable+0x780；在没有 bf16 ALU 的代际上发出"Currently seg scan add for bf16 is only supported"并失败。min/max完全没有 half arm；tpu_{min,max}_seg_scan2xN作为 ops 存在但没有 codegen。- bf16 segmented ISA op 是
SegmentedAddScanBf16PartialSumBf16（oneof0x2e）。 gfc/glc proto pool 没有PartialSumF32segmented form。 - CSR row-offsets 是 segment boundaries。
row_pointers→concatenated_csr_pointers→ 逐 minibatch 切片 →SegmentedScanOpoperand[1]。gathered、gain-scaled embedding rows 是operand[0]。每个 segment 最后一个 lane 上的 inclusive segmented sum 就是 summed embedding。 - intrinsic 是 inclusive。 Exclusive segmented scans 由 front-end 从 inclusive result 合成；硬件原语不暴露 exclusive form。
- 未映射 / LOW。
CreateDynamicSliceCsr的 literal index arithmetic（来自num_minibatches+GetPaddedRowCount的 start/stride）；exclusive-from-inclusive front-end rewrite；最终 packed ISA word 中的 per-port bit positions（emitter slots+0x1c..+0x34已知；encode bit-ranges 由 VEX Mask/Dest-Port/Sub-Opcode 负责）；是否有任何后续代际暴露Bf16PartialSumF32segmented variant。
相关组件
| Name | Relationship |
|---|---|
SegmentedScanOpLowering (0x13589d40) | reduction × dtype switch；发射带 bf16 +0x780 gate 的 tpu_*_seg_scan* family |
SegmentedScanOp::build/create (0x145fd4a0 / 0x145fd5a0) | 绑定 (data, segment) — segment-boundary operand order |
FindAndEmitToUnusedPort (0x13ab2aa0) | 将 segment-id 路由到空闲 V read port 的 7-port greedy allocator |
XlaSparseDenseMatmulWithCsrInputOp::Compile (0xe650800) | CSR row_pointers 成为 segment-id 的 HLO front-end |
MinibatchingDecomposition::CreateDynamicSliceCsr (0x13489ea0) | 将 concatenated_csr_pointers 切成逐 minibatch segment vector |
PackedOperandsLowering ScanOpLowering<SegmentedScanOp> (0x135f3000) | 在最终 lowering 前构建 dialect SegmentedScanOp（unpack → create → pack） |
交叉引用
- Scan Datapath — plain
ScanOplowering、in-scan M-register mask（proto+0x38）、两个 M-register bands、post-scanVectorSelect，以及本页 segmented variant 有意省略的i1/mprefixcount path。 - Segmented Add-Scan —
SegmentedAddScanISA operand frame、2xN/halfform 所依赖的VpackFormatdtype-attribute capability matrix，以及完整的 18 项 segmented-scan proto oneof case map（add_f32=0x23，bf16=0x2e）。 - Embedding Minibatching — minibatching decomposition，以及 segment-id 被切片来源
concatenated_csr_pointers的 provenance。 - VEX Mask / Dest-Port / Sub-Opcode — emitter slots
+0x1c..+0x34和 maskproto+0x38编码到的 bundle bit positions；sub-opcode map。 - VectorExtended (VEX) — segmented scan 发射到的 scan/sort/reduce slot；VEX opcode roster 和 V read ports。
- M-Register Predicate Word (M0–M31) — 独立
proto+0x38mask 字段选择的 predicate word，与 segment operand 不同。 - VectorLoad Slot — gather 本 scan 所归约 embedding rows 的 read-side slot。
- TEC Vector Opcode Enumeration — VEX/
VectorAluopcode roster 和 opcode-recovery model。 - SparseCore 概览 — 三类 SC 引擎以及 TEC vector segmented-scan datapath 的位置。
- Binary:
extracted/libtpu-0.0.40-cp314-cp314-manylinux_2_31_x86_64/libtpu/libtpu.so(build-id89edbbe81c5b328a958fe628a9f2207d) - 索引项： Part IX — SparseCore & BarnaCore / SparseCore datapath (embeddings) — 返回索引