RaggedDot 与卷积几何降级

本页中的所有地址均适用于 libtpu.so v0.0.40（libtpu-0.0.40-cp314，build-id 89edbbe81c5b328a958fe628a9f2207d，781,691,048 B，未 stripped，ELF x86-64，libtpu_lts_20260413_b_RC00）。其他版本会有所不同。所有结论均来自对二进制文件的静态分析。

摘要

XLA 的 kRaggedDot 是一种分组 matmul：一个收缩，其收缩轴（或非收缩轴）被划分为由 group_sizes 操作数描述的一组不规则、可变长度的分组。MXU 没有用于不规则收缩的指令，因此 TPU 编译器会把它降级为硬件已经知道如何发射的形式：一个带窗口的卷积，其空间窗口沿不规则轴移动，再加上一个布尔迭代掩码，用于将稠密卷积不可避免会计算出的跨组乘积清零。执行此操作的 pass 是 xla::jellyfish::RaggedDotExpander，它是位于 platforms/xla/service/jellyfish/ragged_dot_expander.cc 中的 HLO 级 pass（该源路径逐字嵌入在每个错误位置）。

该 expander 的核心数据结构是 RaggedConvSpec（xla::jellyfish::(anonymous namespace)::RaggedConvSpec），由 FromRaggedDot 生成。它是从源指令的 RaggedDotDimensionNumbers 重建出的卷积几何：一个 ConvolutionDimensionNumbers、一个 Window、conv 输出 Shape、feature/batch 分组计数、一个 PrecisionConfig，以及一个单字节收缩模式选择器。随后展开有两个正交选择。第一个是收缩模式（RaggedConvContractionMode，两个分支：reduce 与 dynamic_slice），它决定如何把带掩码的窗口折叠到稠密输出中。第二个是是否附加迭代掩码（xla_tpu_impure_use_iteration_mask 开关，在 Tc-version ≥ 3 时 AUTO=ON），以及在 megacore 部件上是否把掩码拆分到两个核心。expander 运行后，发射出的 kConvolution 会流经普通的 conv → LLO 路径，在那里 SpatialMajorConvolution 根据窗口几何选择一个 ConvolutionLoweringStrategy。

本页记录重新实现者必须复现的三项内容：RaggedConvSpec 几何与 PipelineWindowSpec 窗口边界解析；两个收缩模式分支（带 MaskAggregatorConfig + CreateReduce/CreateAddComputation 的 reduce 分支，以及带 DynamicSliceMaskedConv + update_into 累加器的 dynamic_slice 分支）；以及 Convolve → LLO 的 SpatialMajorConvolution 发射。配套的开关/门控机制（三个 Should…IterationMask 谓词与 impure flag 管线）在此概述，其完整归属在 impure-flags 页面。

对于重新实现，契约是：

RaggedConvSpec 布局 — 它携带哪些 conv 几何字段、位于哪些字节偏移，以及 FromRaggedDot 如何从 RaggedDotDimensionNumbers 推导它们。
窗口模型 — 从 4 元素字符串向量解析出的 g/m/k/n PipelineWindowSpec，以及它如何参数化卷积窗口与 LLO 降级策略。
两个收缩模式分支 — reduce（mask-aggregator + reduce-add）与 dynamic_slice（DynamicSliceMaskedConv + scatter-accumulate），由 RaggedConvSpec 的模式字节选择。
迭代掩码 — Iota + Broadcast + Compare(≥)·Compare(<) + And 布尔掩码，由 kSelect(mask, conv, 0) 三元操作消费，以及何时附加、何时不附加。
LLO 交接 — expander 发射的 kConvolution，以及 SpatialMajorConvolution 如何把窗口转化为 ConvolutionLoweringStrategy。


Pass	`xla::jellyfish::RaggedDotExpander`（HLO pass）
Pass 对象	`0x30` 字节；`+0x08` `use_iteration_mask` bool，`+0x09` `RaggedConvContractionMode`，`+0x10`/`+0x20` `window_bounds` `vector<string>`
Pass 入口	`RunImpl` `0x10fae060`
Pass 接线	`PostMainFusionHloOptimize` `0x10966560`（AddPass 调用点在 `0x109673b3`）→ `AddPass<…>` `0x1096d2e0`（`make_unique` `0x1096e360`）
几何构建器	`FromRaggedDot`（错误 lambda `$_0` `0x10fb2160`）；`ExpandShape` `0x10fb2360`
Expander 主体	`ExpandRaggedDot` `0x10fafa20`
掩码构建器	`CreateOutputMask` `0x10fb2900`
Conv+select 融合	`CreateConvolutionSelectFusion` `0x10fb31e0`；`DynamicSliceMaskedConv` `0x10fb6a00`（dynamic_slice 分支）
LLO 发射器	`xla::jellyfish::SpatialMajorConvolution`（`SetLoweringStrategy` `0x13167e40`）
源文件	`platforms/xla/service/jellyfish/ragged_dot_expander.cc`
IR 层级	HLO（expander）；MLO/LLO（`SpatialMajorConvolution`）

Pass 接线与成员布局

目的

RaggedDotExpander 在主融合 pass 之后，于 PostMainFusionHloOptimize（函数入口 0x10966560，AddPass 调用点在 0x109673b3）中通过 HloPassPipeline::AddPass<RaggedDotExpander, bool, RaggedConvContractionMode, vector<string>, Target&>（0x1096d2e0）构造并添加。三个 impure 编译开关馈入构造函数；它们在流水线构建时读取一次，而不是逐指令读取。

入口点

text

PostMainFusionHloOptimize          0x10966560  (AddPass block at 0x109673b3)
  ├─ RaggedDotExpanderShouldUseIterationMask  0x1d6b5d60  ── use_iteration_mask arg
  ├─ FLAGS_xla_tpu_impure_contract_ragged_conv_with  0x223a7cf8  ── RaggedConvContractionMode arg
  ├─ FLAGS_xla_tpu_impure_ragged_dot_window_bounds   0x223a7d58  ── vector<string> g/m/k/n arg
  └─ AddPass<RaggedDotExpander,…>    0x1096d2e0
       └─ make_unique<RaggedDotExpander>  0x1096e360  (0x30 bytes)
            ├─ +0x08  use_iteration_mask  (bool)
            ├─ +0x09  contraction mode    (RaggedConvContractionMode)
            └─ +0x10 / +0x20  window_bounds (vector<string>)
```text

### 成员布局

| 字段 | 偏移 | 类型 | 含义 |
|---|---|---|---|
| vtable | `+0x00` | ptr | `HloModulePass` vtable |
| `use_iteration_mask` | `+0x08` | bool | 门控带掩码降级；在 ctor 中由 `mov %al,0x8` 写入 |
| 收缩模式 | `+0x09` | enum（1 字节） | `reduce`=0 / `dynamic_slice`=1；由 `mov %cl,0x9` 写入 |
| `window_bounds` | `+0x10`/`+0x20` | `vector<string>` | 4 个字符串 `{g,m,k,n}` |

> **注意 —** 这三个开关都是 *impure*（在编译时读取，不是编译环境 proto 的一部分）。`use_iteration_mask` 与 `enable_masked_fusion_iteration_skipper` 是 `AutoOr<bool>`，直接从其 `FlagImpl+0x58` 缓存读取；`contract_ragged_conv_with` 是一个 `RaggedConvContractionMode` `FixedOptionSetFlag`；`ragged_dot_window_bounds` 是一个 `vector<string>`。完整的 impure-flag 机制（AUTO 极性、`FlagImpl` 读取路径）归属于 [Registry-Mediated Flags](../config/registry-mediated-flags.md)。

### 三个迭代掩码门控

存在两个不同的门控函数，因为掩码会在两个流水线层级被咨询：HLO expander 与 LLO conv emitter；它们读取不同谓词。三者都首先基于 TPU Tc-version ≥ 3 门控（`mov 0x8(view),%rax; cmpl $0x3,(%rax); jl ret-false`）。

| 门控 | 地址 | 谓词 | 消费者 |
|---|---|---|---|
| `RaggedDotExpanderShouldUseIterationMask` | `0x1d6b5d60` | v≥3 AND `use_iteration_mask`（AUTO=ON） | `PostMainFusionHloOptimize` `0x10966560`（`+0x08` ctor 参数） |
| `ShouldUseIterationMask` | `0x1d6b5dc0` | v≥3 AND（`use_iteration_mask` OR `enable_masked_fusion_iteration_skipper`） | `SpatialMajorConvolution` ctor / Emit（LLO 层级） |
| `ShouldEnableMaskedFusionIterationSkipper` | `0x1d6b5d20` | v≥3 AND 普通 bool skipper | （上方的 skipper 析取项） |

expander 门控（`0x1d6b5d60`）只咨询 `use_iteration_mask`；LLO 门控（`0x1d6b5dc0`）还遵守 masked-fusion-iteration-skipper。AUTO=ON 极性是 `AutoOr` 惯用法：读取缓存字，并且结果为 `(word & 0x101) != 0x100`，即除非 flag 被显式设置为 present-and-false，否则为 true。

```c
function ShouldUseIterationMask(env, topology):     // 0x1d6b5dc0
    if topology.tc_version < 3: return false         // mov 0x8(a2); cmpl $3; jl
    // disjunct 1: use_iteration_mask (AutoOr<bool>, AUTO=ON)
    if FLAGS_use_iteration_mask present:
        word = ~ReadOneWord(FLAGS_xla_tpu_impure_use_iteration_mask)
        if (word & 0x101) == 0: return true
    // disjunct 2: the masked-fusion-iteration-skipper (plain bool)
    if topology.tc_version >= 3:
        if ReadOneBool(FLAGS_xla_tpu_impure_enable_masked_fusion_iteration_skipper):
            return true
    // final: AUTO resolves to ON
    word = ReadOneWord(FLAGS_xla_tpu_impure_use_iteration_mask)
    return (word & 0x101) != 0x100

RaggedConvSpec 几何

目的

RaggedConvSpec 是 RaggedDot 会变成的卷积。它是 ragged_dot_expander.cc 中 anonymous-namespace 的 struct，由 FromRaggedDot(const HloInstruction*, RaggedConvContractionMode, const Target&) 构建。它携带 HloInstruction::CreateConvolve 所需的一切，以及模式选择器。其字段偏移直接从 CreateConvolutionSelectFusion（0x10fb31e0）里的 CreateConvolve 调用恢复，其中每个 conv 参数都从 spec 指针（a7/v108）的固定偏移加载。

布局

下列偏移从单个 CreateConvolve 调用点（0x10fb31e0，约第 1288 行）及其周围加载的反汇编读出。spec 指针以 int64*（v108）持有，因此反编译器把 conv 子对象指针渲染为 qword 索引 — v108 + 6、v108 + 24、v108 + 29、v108 + 72 — 它们对应字节偏移 8 × 索引。收缩模式选择器是唯一例外：它通过字节访问 *(int8*)(a7 + 24) 读取，也就是真正的字节偏移 24，不是 qword 索引 24。调用是 CreateConvolve(builder, lhs, rhs, feature_group_count, batch_group_count, &window, &dim_numbers, &precision_config, &sparsity, 0)，其中 v108+24、v108+6、v108+72 作为 window/dim-numbers/precision 参数（第 1288 行），并且 Shape::Shape(&v417, v108+29) 作为输出 shape（第 1272 行，也即第 453 行的 a7+232）：

字段	字节偏移	反编译中的 qword 索引	类型	在 `CreateConvolve` 中的角色
收缩模式字节	`+0x18` (24)	—（字节访问 `(int8)(a7+24)`）	u8	`0`=reduce 分支，`1`=dynamic_slice 分支；测试后选择 fold
`dim_numbers`	`+0x30` (48)	`v108 + 6`	`ConvolutionDimensionNumbers`	conv dim numbers 参数
`window`	`+0xC0` (192)	`v108 + 24`	`Window`	window 参数
conv `Shape`	`+0xE8` (232)	`v108 + 29`	`Shape`	输出/conv shape（`Shape::Shape(&v417, a7+232)`）
`feature_group_count`	`+0x230` (560)	`v108[70]`	i64	conv feature-group 参数
`batch_group_count`	`+0x238` (568)	`v108[71]`	i64	conv batch-group 参数
`precision_config`	`+0x240` (576)	`v108 + 72`	`PrecisionConfig`	precision 参数

易错点 — 收缩模式字节位于字节 24，但 Window proto 位于 qword 索引 24，即字节 192，不是字节 24。反编译器把两者都渲染为字面量 24，但模式是单字节读取（*(int8*)(a7+24)），而窗口是按 qword 索引的子对象指针（v108 + 24 → a7 + 192）；它们是相距约 168 字节的不同字段。把每个 qword 索引乘以 8 才能恢复字节偏移。在加载无歧义的地方，字节偏移列为 HIGH/MEDIUM 置信度。

RunImpl 通过其析构函数确认了 spec 的组件集合：每次展开尝试之后，它都会拆除一个 PrecisionConfig、一个 Shape、一个 Window 和一个 ConvolutionDimensionNumbers（0x10fae060，约第 1408-1419 行）— 正是上方四个 conv 子对象。

FromRaggedDot

FromRaggedDot 验证 RaggedDotDimensionNumbers 并组装 spec；唯一单独符号化的部分是其错误 lambda $_0（0x10fb2160），它格式化 "Failed to create RaggedConvSpec from ragged_dot instruction %s: %s, got %lld"（MakeErrorImpl<3>，源第 145 行）。驱动该 lambda 的验证内联在 RunImpl 中：

function RunImpl(module, …):                         // 0x10fae060
    for comp in module.MakeComputationPostOrder():    // 0x10fae060+250
      for inst in comp where inst.opcode == kRaggedDot:
        // a ragged BATCH dim must already be gone
        CHECK(ragged_dot_mode != RaggedDotMode::kBatch)        // "Ragged dot with a ragged batch dim
                                                               //  should not reach the expander."
        // exactly-one-of dim validation (each failure -> FromRaggedDot $_0 error)
        CHECK(num_contracting_dims == 1)              // "number of contracting dimensions should be 1"
        CHECK(num_lhs_noncontracting == 1)            // "number of lhs non-contracting dimensions should be 1"
        CHECK(num_rhs_noncontracting == 1)            // "number of rhs non-contracting dimensions should be 1"
        spec = FromRaggedDot(inst, mode, target)      // build RaggedConvSpec
        if use_iteration_mask:                         // member +0x8 == 1
            CHECK(window_bounds.size() == 4)          // "window_bounds.size() == 4"
            pws.g = SimpleAtoi(window_bounds[0])       // safe_strto64_base(…,10)
            pws.m = SimpleAtoi(window_bounds[1])
            pws.k = SimpleAtoi(window_bounds[2])
            pws.n = SimpleAtoi(window_bounds[3])
        ExpandRaggedDot(inst, spec, use_iteration_mask,
                        optional<PipelineWindowSpec>(pws), nullopt, target)   // 0x10fafa20
```text

> **怪癖 —** `kBatch` ragged 模式是硬 `CHECK`，不是优雅降级。ragged *batch* 维度预期已被更早的 pass 消除；如果有一个存活到 expander，它会中止。只有 ragged *contracting* 与 ragged *non-contracting* 模式会到达 `FromRaggedDot`，并且各自都要求其对应维度类别恰好有一个维度。

### PipelineWindowSpec 窗口边界

当设置 `use_iteration_mask` 时，expander 通过 `absl::SimpleAtoi` / `safe_strto64_base(…, 10)` 将四个 `window_bounds` 字符串解析为 `PipelineWindowSpec {g, m, k, n}`。映射是位置固定的（每个解析目标都在其 `CHECK` 字符串中命名，位于 `0x10fae060+1295…1359`）：

| 索引 | 字段 | 源 CHECK 字符串 |
|---|---|---|
| `window_bounds[0]` | `pipeline_window_spec.g` | `"absl::SimpleAtoi(window_bounds[0], &pipeline_window_spec.g)"` |
| `window_bounds[1]` | `pipeline_window_spec.m` | `"…window_bounds[1], &pipeline_window_spec.m)"` |
| `window_bounds[2]` | `pipeline_window_spec.k` | `"…window_bounds[2], &pipeline_window_spec.k)"` |
| `window_bounds[3]` | `pipeline_window_spec.n` | `"…window_bounds[3], &pipeline_window_spec.n)"` |

`g/m/k/n` 是 grouped-matmul 轴（groups、lhs-non-contracting M、contracting K、rhs-non-contracting N）。`PipelineWindowSpec` 是*流水线*窗口，不同于 `RaggedConvSpec` 内部的 conv `Window` proto；它作为 `optional` 参数传入 `ExpandRaggedDot`，再传到 `CreateConvolutionSelectFusion`。它只在带掩码路径上填充；无掩码路径传递 `nullopt`。

---

## 收缩模式分支

`RaggedConvContractionMode` 是一个 `FixedOptionSetFlag`，在 `GetRaggedConvContractionModeParser`（`0x1db15340`）中注册，恰好有两个选项：

```text
"reduce"        -> 0      // mask-aggregator + reduce-add fold
"dynamic_slice" -> 1      // DynamicSliceMaskedConv + scatter-accumulate

模式字节随 RaggedConvSpec（+0x18）携带，并在整个 CreateConvolutionSelectFusion 中以 *(a7+24) 测试，用于在两个 fold 之间切换。两个分支共享相同的前半部分：对带窗口的操作数执行稠密 kConvolution，然后应用逐窗口掩码。它们的区别在于如何把带掩码窗口组合成稠密输出。

Reduce 分支（`reduce`，模式 0）

reduce 分支构建一个融合 computation：执行卷积、按迭代掩码 select、广播零值，并跨窗口执行 reduce-add。从 CreateConvolutionSelectFusion（0x10fb31e0）以及 ExpandRaggedDot 中的 MaskAggregatorConfig 设置恢复如下：

// inside CreateConvolutionSelectFusion, *(spec+24) == 0  (reduce)
function ReduceArmFusionBody(spec, mask, …):           // 0x10fb31e0
    kernel       = Parameter(0, "kernel")
    activations  = Parameter(1, "activations")
    output_mask  = Parameter(2, "output_mask")
    conv  = CreateConvolve(builder, kernel, activations,
                           spec.feature_group_count, spec.batch_group_count,
                           &spec.window, &spec.dim_numbers, &spec.precision, &sparsity, 0)
    zero  = CreateBroadcast(CreateConstant(Zero(elem_type)))   // 0
    sel   = CreateTernary(108 /*kSelect*/, output_mask, conv, zero)  // select(mask, conv, 0)
    addc  = CreateAddComputation(elem_shape)                   // scalar add reducer
    out   = CreateReduce(sel, init=zero, dims={window}, addc)  // sum masked windows
    return Build(out)                                          // wrapped in a kFusion
```text

`MaskAggregatorConfig`（一个 `xla::jellyfish` proto，在 `ExpandRaggedDot` 的 `0x10fafa20` 中以 arena `DefaultConstruct`）参数化聚合。它的两个字段被直接写入：`+44` 处的 flag 字节设置为 `0`，`+48` 处的打包二元组设置为 `0x200000000`（即以两个 32 位半字表示的 `{0, 2}`），并设置 presence bits `(+16) |= 0xE`。同一设置块还配置了一个 `WindowConfig`（`+184 = 1`，presence `(+17) |= 0x20`）。在 megacore 部件上，spec 还会进一步拆分（见下文）。

> **注意 —** `MaskAggregatorConfig` 是迭代掩码在 SparseCore/TensorCore 层级的实现：它告诉发射出的 kernel 如何跨 conv 窗口聚合逐组掩码。只恢复了它的两个已写字段（`+44`、`+48 = 0x200000000`）；完整 proto schema 未解码（除观测到的写入外，对字段语义的置信度为 LOW）。

### Dynamic-Slice 分支（`dynamic_slice`，模式 1）

当 `*(spec+24) == 1` 时，fold 委托给 `xla::jellyfish::DynamicSliceMaskedConv`（`0x10fb6a00`，从 `CreateConvolutionSelectFusion` 约第 2003 行调用）。该分支额外添加两个融合参数：`group_starts`（`Parameter(3)`）和 `update_into`（`Parameter(4)`），并将每个分组的带掩码 conv 结果按该分组的起始偏移 scatter 到运行中的累加器，而不是对窗口执行 reduce-sum：

```c
// inside CreateConvolutionSelectFusion, *(spec+24) == 1  (dynamic_slice)
function DynamicSliceArm(spec, mask, group_starts, update_into, …):    // 0x10fb31e0 -> DynamicSliceMaskedConv 0x10fb6a00
    CHECK(group_starts_param != nullptr)   // "group_starts_param != nullptr"
    CHECK(update_into_param  != nullptr)   // "update_into_param != nullptr"
    // build the conv + select(mask, conv, 0) as in the reduce arm, then:
    DynamicSliceMaskedConv(builder, spec, group_starts, update_into, …)
    // scatters the masked conv into update_into at the per-group start offset

对于正确的不规则收缩，这两个分支会产生数值等价的结果；dynamic_slice 通过把每个分组的贡献写入原位来避免完整窗口 reduce，这在分组较大且稀疏时更可取。选择完全由 xla_tpu_impure_contract_ragged_conv_with 开关决定。

ConvDimNumbers

ConvolutionDimensionNumbers 是 spec 的 +0x06 字段，它使 ragged 收缩轴成为卷积空间维度。它从 ExpandRaggedDot 中源 RaggedDotDimensionNumbers 重建（RaggedDotDimensionNumbers::RaggedDotDimensionNumbers((…)v313, 0, ragged_dot_dimension_numbers(inst))，0x10fafa20+1063），并在 CreateConvolutionSelectFusion 中多次 copy-construct（0x10fb31e0+1631/1917/1942）。窗口的 ragged 维度会针对 kernel/output 窗口边界做健全性检查：CHECK(conv_window_config.output_window_bounds_size() == 3)、CHECK(conv_window_config.kernel_window_bounds_size() == 3) 和 CHECK(ragged_dim_window_bound <= max_ragged_dim_window_bound)。

ExpandRaggedDot — 子图

目的

ExpandRaggedDot（0x10fafa20）是用卷积子图替换 kRaggedDot 指令的例程。它按顺序构建：分组边界机制（来自 group_sizes）、带窗口且填充后的操作数、迭代掩码，以及 conv+select 融合。整个过程以 HloComputation::ReplaceInstruction(inst, fused_conv) 结束。

算法

function ExpandRaggedDot(inst, spec, use_iteration_mask, pws, dyn, target):   // 0x10fafa20
    // 0. optional control dependency on operand 2 (group_sizes) via AfterAll+AddDependency
    gs = inst.mutable_operand(2)                       // group_sizes
    if has_control_dep: gs = AddDependency(gs, AfterAll())

    // 1. group_ends  = reduce-window prefix over group_sizes
    CHECK(gs.shape.rank == 1)                           // "group_sizes should be rank 1"
    gs2  = Bitcast(gs, AppendMajorDimension(gs.shape))  // conv layout
    zero = CreateConstant(LiteralUtil::Zero(elem))
    addc = CreateAddComputation(...)
    group_ends = CreateReduceWindow(gs2, zero, window, addc)   // running group offsets
    group_ends = AddInstruction(... UniquifyName "group_ends")

    // 2. group boundaries fusion (group_starts from group_ends)
    CHECK(group_ends.shape.rank == 2)                   // "group_ends should be rank 2"
    bc   = Bitcast(StripShape(group_ends))
    off  = CreateConstant(LiteralUtil::CreateR1<int>({...}))   // group offsets literal
    // pad_add_fusion: Pad(group_ends) AND Pad(slice) -> kAnd -> kFusion "pad_add_fusion"
    pad1 = CreatePad(group_ends, zero, cfg)             // "group boundaries fusion"
    pad2 = CreatePad(slice,      zero, cfg)
    bnd  = CreateBinary(kAnd, pad1, pad2)               // "group_starts_one_longer should be rank 1"
    bnd_fusion = CreateFusion(Build(bnd))               // -> group starts/ends

    // 3. iteration mask(s)
    group_starts = Slice(bnd_fusion, ...)               // group_starts from boundaries
    out_shape    = AppendMinorDimension(operand0.shape) // conv output layout
    mask  = CreateOutputMask(group_ends_inst, group_starts, out_shape, target)   // 0x10fb2900
    if use_iteration_mask AND Megacore AND num_cores >= 2:
        mask2 = CreateOutputMask(..., AppendMinorDimension x2, ...)              // per-core mask
        // operand0 also wrapped in AfterAll + AddDependency for the second core

    // 4. the conv + select fusion (picks reduce vs dynamic_slice from spec.mode)
    fused = CreateConvolutionSelectFusion(operand0, operand1, mask,
                                          out_shape, spec, pws, target, use_iteration_mask)  // 0x10fb31e0

    // 5. megacore split: tag the MaskAggregator + fusion with the split dim
    if use_iteration_mask AND Megacore AND num_cores >= 2:
        wd = windowing_util::MakeWindowDescription(target, fused.shape)
        if even_split:  AddMatchingMegacoreSplitDimensionsToMaskAggregatorAndFusion(fused, …)
        else:           AddMegacoreSplitDimension(fused, …)  // + MegacoreConfig

    return ReplaceInstruction(inst, fused)
```text

> **怪癖 —** 在拥有 ≥ 2 个核心的 megacore 部件上，expander 会构建**两个**掩码（`CreateOutputMask` 被调用两次，并逐步追加 minor 维度），并且要么调用 `AddMatchingMegacoreSplitDimensionsToMaskAggregatorAndFusion`（当 ragged 维度在核心之间均匀拆分时），要么调用 `AddMegacoreSplitDimension` + 一个 `MegacoreConfig`（否则）。单核心或非 megacore 的重新实现只需要第一个掩码，不需要任何 split helper。拆分决策是 ragged-dim 窗口边界对每核心拆分的整除性测试。

### 迭代掩码 — CreateOutputMask

`CreateOutputMask`（`0x10fb2900`）构建用于将跨组乘积清零的布尔掩码。它是该降级的结构核心：沿带窗口 ragged 轴进行的稠密卷积会计算完整的块笛卡尔积，而掩码只保留每个输出位置上的 `[group_start, group_end)` 带。

```c
function CreateOutputMask(group_ends, group_starts, out_shape, target):   // 0x10fb2900
    body builds a kFusion computation:
      stripped_ends = Parameter(0, "stripped_group_ends")    // upper bounds
      group_starts  = Parameter(1, "group_starts")           // lower bounds
      lo   = CreateBroadcast(group_starts, out_shape)
      iota = CreateIota(out_shape, iota_dim)                 // output-position index
      ge   = CreateCompare(lo, iota, 4 /*kGe*/)              // iota >= group_start
      hi   = CreateBroadcast(stripped_ends, out_shape)
      lt   = CreateCompare(hi, iota, 5 /*kLt*/)              // iota <  group_end
      mask = CreateBinary(13 /*kAnd*/, ge, lt)               // in-band <=> in this group
    return CreateFusion(Build(mask))                          // boolean mask, one kFusion

易错点 — 比较方向是下界使用 kGe（opcode 4）、上界使用 kLt（opcode 5），并与同一个输出位置 kIota 比较。带区间是半开区间 [group_start, group_end)。随后掩码被 kAnd（opcode 13）合并，并包裹在自己的 kFusion 中，使布尔计算独立于 conv 融合。半开约定或比较极性出错，会悄悄地对每个分组的边界元素双计或漏计。

掩码馈入 CreateConvolutionSelectFusion，其中 opcode 为 108（kSelect）的 CreateTernary 计算 select(mask, conv, broadcast_zero) — 带掩码的 conv 结果。reduce 分支随后对它执行 reduce-add；dynamic_slice 分支则 scatter 它。

Convolve → LLO：SpatialMajorConvolution

目的

expander 留下一个普通的 kConvolution（位于 kFusion 内部）。它流经标准 conv → MLO/LLO 路径。在 LLO 层级，xla::jellyfish::SpatialMajorConvolution 是把该卷积转化为 tiled、spatially-major matmul-accumulate 循环的发射器。它会通过 LLO 门控 ShouldUseIterationMask（0x1d6b5dc0）重新咨询迭代掩码谓词，这也是为什么该门控还遵守 masked-fusion-iteration-skipper。

降级策略选择

SpatialMajorConvolution 在对象偏移 +0x882（2178）处持有一个 ConvolutionLoweringStrategy。它有两种设置方式：

function SpatialMajorConvolution::SetLoweringStrategy(strategy):     // 0x13167e40
    this[+0x882] = strategy                          // store the 24-byte strategy struct

function SpatialMajorConvolution::UpdateLoweringStrategyWithWindowInfo(   // 0x13167e80
        ragged_dim_bound, lhs_bound, rhs_bound, is_ragged):
    s = convolution_util::GetConvolutionLoweringStrategy(    // window geometry -> strategy
            this.window_field, this.spatial_field,
            ragged_dim_bound, /*flag=*/1, is_ragged + 256, …)
    this[+0x882] = s
```text

`UpdateLoweringStrategyWithWindowInfo` 通过 `convolution_util::GetConvolutionLoweringStrategy` 根据窗口边界推导策略 — `g/m/k/n` 窗口几何就是在这里变成具体的 LLO tiling 决策。配套的 `GetLoweringStrategyString`（`0x13167e60`）会将其渲染为诊断信息（`convolution_util::GetLoweringStrategyString(this, this+2178)`）。

### 函数映射

| 函数 | 地址 | 角色 |
|---|---|---|
| `SpatialMajorConvolution::SetLoweringStrategy` | `0x13167e40` | 在 `+0x882` 存储 strategy |
| `SpatialMajorConvolution::GetLoweringStrategyString` | `0x13167e60` | 为诊断渲染 strategy |
| `SpatialMajorConvolution::UpdateLoweringStrategyWithWindowInfo` | `0x13167e80` | 窗口几何 → `GetConvolutionLoweringStrategy` → 存储 |
| `SpatialMajorConvolution::RoundUpWindowBoundToFactorAndCompact` | `0x1315b2a0` | 将 ragged 窗口边界向上取整到硬件因子 |
| `SpatialMajorConvolution::InitFromFusion` | `0x13155fc0` | 将发射器绑定到 expander 的 `kFusion` |
| `SpatialMajorConvolution::InitFusionEmitters` | `0x13173b60` | 每操作数窗口发射器 |
| `SpatialMajorConvolution::MatrixMultiplyAccumulate` | `0x131792e0` | MXU MMA 内部发射 |
| `SpatialMajorConvolution::EmitZeroByteCase` / `EmitZeroElementCases` | `0x13178100` / `0x13178020` | 退化窗口短路 |
| `SpatialMajorConvolution::PopulateNestedOutputFusions` | `0x13155d40` | 接线 select/reduce 子融合 |

> **注意 —** `MatrixMultiplyAccumulate`（`0x131792e0`）是通往与 dense dot/conv 共享的 MXU emitter 的桥梁。本页止于策略决策；它分派进入的 tile-cost comparator 与 `EmitFunctorEnum` 是 [Dot / Conv → MXU Lowering](dot-conv-mxu-lowering.md) 的主题。

---

## 函数映射

| 函数 | 地址 | 角色 |
|---|---|---|
| `RaggedDotExpander::RunImpl` | `0x10fae060` | 逐指令验证 + 窗口解析 + 分派 |
| `ExpandRaggedDot` | `0x10fafa20` | 构建 conv 子图，替换 `kRaggedDot` |
| `FromRaggedDot`（错误 lambda `$_0`） | `0x10fb2160` | `RaggedConvSpec` 构建失败格式化器 |
| `ExpandShape` | `0x10fb2360` | 将 conv 空间维度追加到操作数 shape |
| `CreateOutputMask` | `0x10fb2900` | `Iota + Broadcast + Compare(≥)·Compare(<) + And` 掩码 |
| `CreateConvolutionSelectFusion` | `0x10fb31e0` | `Convolve + Select(mask,·,0)` + reduce/dynamic_slice fold |
| `DynamicSliceMaskedConv` | `0x10fb6a00` | `dynamic_slice` 分支 scatter-accumulate |
| `GetRaggedConvContractionModeParser` | `0x1db15340` | `FixedOptionSetFlag`（`reduce`/`dynamic_slice`） |
| `RaggedDotExpanderShouldUseIterationMask` | `0x1d6b5d60` | expander 级掩码门控 |
| `ShouldUseIterationMask` | `0x1d6b5dc0` | LLO 级掩码门控 |
| `ShouldEnableMaskedFusionIterationSkipper` | `0x1d6b5d20` | skipper 析取项 |
| `make_unique<RaggedDotExpander>` | `0x1096e360` | ctor + 成员布局 |
| `AddPass<RaggedDotExpander,…>` | `0x1096d2e0` | 流水线接线 |

---

## 诊断字符串

全部由 `platforms/xla/service/jellyfish/ragged_dot_expander.cc` 发出。这些是重新实现必须遵守的核心断言（也是调试工程师会 grep 的内容）。

| 字符串 | 触发时机 | 严重性 |
|---|---|---|
| `Ragged dot with a ragged batch dim should not reach the expander.` | `ragged_dot_mode == kBatch` 到达 `RunImpl` | FATAL CHECK |
| `number of contracting dimensions should be 1` | `FromRaggedDot` 中的 dim 验证 | error Status |
| `number of lhs non-contracting dimensions should be 1` | dim 验证 | error Status |
| `number of rhs non-contracting dimensions should be 1` | dim 验证 | error Status |
| `Failed to create RaggedConvSpec from ragged_dot instruction %s: %s, got %lld` | `FromRaggedDot` `$_0` | error Status |
| `window_bounds.size() == 4` | 带掩码路径，畸形 `ragged_dot_window_bounds` | FATAL CHECK |
| `absl::SimpleAtoi(window_bounds[{0..3}], &pipeline_window_spec.{g,m,k,n})` | 非数字窗口边界 | error Status |
| `group_sizes should be rank 1` | `group_sizes` 操作数 shape | error Status |
| `group_ends should be rank 2` | reduce-window 输出 shape | error Status |
| `group_starts_one_longer should be rank 1` | 边界融合 shape | error Status |
| `conv_window_config.output_window_bounds_size() == 3` | conv 窗口组装 | FATAL CHECK |
| `conv_window_config.kernel_window_bounds_size() == 3` | conv 窗口组装 | FATAL CHECK |
| `ragged_dim_window_bound <= max_ragged_dim_window_bound` | 窗口边界 clamp | FATAL CHECK |
| `group_starts_param != nullptr` / `update_into_param != nullptr` | `dynamic_slice` 分支前置条件 | FATAL CHECK |

---

## 相关组件

| 名称 | 关系 |
|---|---|
| `RaggedDotExpander` | 本文档记录的 HLO pass |
| `SpatialMajorConvolution` | 降级 expander 产生的 `kConvolution` 的 LLO 发射器 |
| `MaskAggregatorConfig` / `WindowConfig` / `MegacoreConfig` | 由 `ExpandRaggedDot` 写入的 `jellyfish` backend-config protos |
| `RaggedConvContractionMode` | `reduce` / `dynamic_slice` `FixedOptionSetFlag` |
| MXU dot/conv emitter | `MatrixMultiplyAccumulate` 将带窗口 conv 交给共享 MXU 路径 |

---

## 交叉引用

- [编译器概览](overview.md) — `RaggedDotExpander` 在 post-fusion HLO 流水线中的位置
- [Dot / Conv → MXU 降级](dot-conv-mxu-lowering.md) — `MatrixMultiplyAccumulate` 分派进入的 tile-cost comparator / `EmitFunctorEnum`
- [TPU → LLO ODS](tpu-to-llo-ods.md) — `SpatialMajorConvolution` 面向的 LLO op 表面
- [融合模式](fusion-patterns.md) — 掩码与 conv+select 融合遵循的 `kFusion` 包装约定
- [Custom-Call 降级](custom-call-lowering.md) — 对无原生 MXU 形式 ops 的同级 Part V 降级路径
- [Registry-Mediated Flags](../config/registry-mediated-flags.md) — `use_iteration_mask` 与 `contract_ragged_conv_with` 背后的 impure/`AutoOr` flag 读取路径
- [返回索引](../index.md) — Part V — 编译器：降级与优化 Pass

RaggedDot 与卷积几何降级 ​

摘要 ​

Pass 接线与成员布局 ​

目的 ​

入口点 ​

RaggedConvSpec 几何 ​

目的 ​

布局 ​

FromRaggedDot ​

Reduce 分支（reduce，模式 0） ​

ConvDimNumbers ​

ExpandRaggedDot — 子图 ​

目的 ​

算法 ​

Convolve → LLO：SpatialMajorConvolution ​

目的 ​

降级策略选择 ​

RaggedDot 与卷积几何降级

摘要

Pass 接线与成员布局

目的

入口点

RaggedConvSpec 几何

目的

布局

FromRaggedDot

Reduce 分支（`reduce`，模式 0）

ConvDimNumbers

ExpandRaggedDot — 子图

目的

算法

Convolve → LLO：SpatialMajorConvolution

目的

降级策略选择