LowerToMlo DMA Bridge-Cast

本页中的所有地址、符号和偏移都适用于 libtpu-0.0.40-cp314 wheel 中的 libtpu.so（build-id 89edbbe81c5b328a958fe628a9f2207d，781,691,048 字节，未 stripped，.text VA == 文件偏移；构建 libtpu_lts_20260413_b_RC00）。其他版本会不同；请把每个 VA 都视为版本绑定。

摘要

createLowerToMloPass (0x1322adc0) 是把 tpu 方言降到 SparseCore 中层级 (Mlo / sparse_core) 方言的 MLIR ModuleOp pass。它是 tpu → LLO ODS 的 SparseCore 对应物（后者是 TensorCore FunctionPass 向 LLO 的下降）。LowerToLLO 会一次性把每个 tpu op 降到一个 LLO 目标，而 LowerToMlo 是一个两阶段 lowering：大多数 tpu op 在这个 pass 中直接降为 sc_tpu.*（sparse_core 方言）/ arith / memref op，但 DMA 和同步族 — tpu.enqueue_dma、tpu.enqueue_indirect_dma、tpu.wait_dma2、tpu.fetch_and_add_sync — 不会在这里 lowering。这些 op 会物化一个带标签的 builtin.unrealized_conversion_cast（即 bridge-cast），把真正的 DMA 发射推迟到下游 SparseCore pass。

本页负责说明该桥接的 LowerToMlo 专属机制：

两阶段 DMA bridge-cast — tpu.enqueue_dma（以及三个同族 DMA/sync op）如何 lowering 为按操作数生成的、标记为 sc.unlowering 的 builtin.unrealized_conversion_cast，源 op 如何标记为 sc.unlowered，以及该 cast 向前携带什么。
块签名转换 — func::FuncOp 动态合法性谓词如何检测 SparseCore sequencer 函数，并强制其签名经过 TypeConverter；以及 UnrealizedConversionCastOp 合法性守卫如何允许 bridge-cast 在传递中穿过 legalizer。
延迟 DMA 物化解析 — 带标签的 bridge-cast 如何在 applyFullConversion 后存活，并随后由 ExpandTiledMemRefsPass 内的 substituteUnloweringConversionCastOp (0x134e73e0) 消费，以发射真正的 SparseCore DMA op。

让 legalization 安全的通用 speculative-apply / rollback 引擎（IRRewrite 动作日志、undoRewrites / resetState / applyRewrites、UnresolvedMaterializationRewrite 记录）不在本页 — 它位于 ConversionPatternRewriter。本页只覆盖 LowerToMlo 桥接、LowerToMlo 合法性谓词以及延迟 DMA 解析。


Pass 工厂	`mlir::tpu::createLowerToMloPass(xla::jellyfish::Target const&, sparse_core::LowerToMloPassContext*)` — `0x1322adc0`
Pass 入口	`mlir::tpu::(anon)::LowerToMloPass::runOnOperation()` — `0x1322b200`
Pass 种类	`ModuleOp` pass (`jaxlib::mlir::Pass<LowerToMloPass, mlir::ModuleOp>`) — 不是 `FunctionPass`
驱动器	`mlir::applyFullConversion(module, MloConversionTarget, frozenPatterns, ConversionConfig)` — 不是 partial
目标方言	`sparse_core`（`Mlo` 中层级）+ `arith`/`memref`/`scf`/`vector`/`cf`/`func`/`math`/`LLVM`
合法方言集合	`MloConversionTarget` ctor `0x13245900`（vtable `0x21903b98`）：11 个方言 `TypeID`
Pattern 集合	`populateTpuToMloConversionPatterns` (`0x1322c920`) — 44 个函数式 `ConversionPattern` lambda + 1 个 catch-all
DMA bridge op	`tpu.enqueue_dma` (`0x13239e00`)、`tpu.enqueue_indirect_dma` (`0x1323a660`)、`tpu.wait_dma2` (`0x1323ae00`)、`tpu.fetch_and_add_sync` (`0x1323b600`)
桥接目标	标记为 `sc.unlowering` 的 `builtin.unrealized_conversion_cast`（按操作数）+ 源 op `sc.unlowered`
第二阶段消费者	`ExpandTiledMemRefsPass` 下的 `substituteUnloweringConversionCastOp` (`0x134e73e0`)
来源溯源	`.rodata` 字符串 `platforms/xla/mosaic/dialect/tpu/transforms/lower_to_mlo.cc`（`EnqueueDMA` 函数体第 405 行）

为什么分两阶段 — 直接 lowering 与桥接的分割

目的

LowerToMlo 处理的每个 tpu op 都落入两个桶之一。先建立这个分割，才能让页面其余部分可读：bridge-cast 不是兜底，而是对一个最终形态必须等 tile expansion 运行后才知道的 op 族进行的有意延迟。

分割

populateTpuToMloConversionPatterns (0x1322c920) 安装 44 个函数式 ConversionPattern lambda（签名 LogicalResult(SrcOp, SrcOpAdaptor, ConversionPatternRewriter&)）。其中 40 个把源 op 直接 lowering 为一个或多个 sc_tpu.*（Mlo / sparse_core 方言）op — 经常是 1:N 展开（例如 tpu.iota → sc_tpu.vlaneseq [+ arith.index_cast]，tpu.device_id → 6 个 op，tpu.delay → 由 sc_tpu.sdelay 组成的 scf.for 嵌套）。这些直接 lowering 是 LowerToLLO 函数体的 SparseCore 类比；TensorCore 侧已在 tpu → LLO ODS 概览，本页不再重新列成表。

剩下的四个 op 是 DMA 和同步族。它们在 LowerToMlo 中 lowering 为空。相反，每个都会发射一个 builtin.unrealized_conversion_cast — bridge-cast — 在转换后的 (Mlo) 类型系统中向前携带原始 op 的操作数，并打上标签供后续 pass 查找。

`tpu` 源 op	LowerToMlo lambda	Lowering 目标	阶段
`tpu.enqueue_dma`	`0x13239e00`	带标签的 `builtin.unrealized_conversion_cast`	桥接
`tpu.enqueue_indirect_dma`	`0x1323a660`	带标签的 `builtin.unrealized_conversion_cast`	桥接
`tpu.wait_dma2`	`0x1323ae00`	带标签的 `builtin.unrealized_conversion_cast`	桥接
`tpu.fetch_and_add_sync`	`0x1323b600`	带标签的 `builtin.unrealized_conversion_cast`	桥接
`tpu.wait_indirect_dma`	`0x1323be00`	`sc_tpu.stream_wait` (`StreamWaitOp`)	直接

注意 — tpu.wait_indirect_dma (0x1323be00) 不是 bridge op — 它直接 lowering，在计算动态大小和 granule 计数（lowering_util::GetDynamicSizes / AssertAlignmentAndGetNumGranules）后发射一个 sparse_core::StreamWaitOp (sc_tpu.stream_wait)。只有上面四个 op 会延迟。这个分割的原因是结构性的：enqueue_dma 的最终 SparseCore 形态（rolled / retiled 传输循环、granule 分解、host-IOVA 与芯片内路由）取决于 tiled memref 布局，而该布局直到 ExpandTiledMemRefsPass 运行后才会解析。LowerToMlo 无法知道形状，所以它把操作数保存在带类型的 cast 后面，让 tile-aware pass 完成工作。消费者是 ExpandTiledMemRefsPass::addPattern<tpu::EnqueueDMAOp> (0x134eef60)、addPattern<tpu::WaitDMA2Op> (0x134ef720)、addPattern<tpu::FetchAndAddSyncOp> (0x134f1240)、addPattern<tpu::EnqueueIndirectDMAOp> (0x134ef340)，这确认了这一点 — 四个 bridge op 正好就是该 pass 重新处理的四个 op。

阶段 1 — DMA Bridge-Cast（`tpu.enqueue_dma` → 带标签 cast）

目的

这是本页的核心：重新实现者必须复现的精确发射形式，这样延迟 op 才能在 full conversion 中存活，并能在下游恢复。tpu.enqueue_dma lambda 函数体 (0x13239e00) 已完整反编译；另外三个 bridge op 遵循同一形态。

算法

该 lambda 遍历 op 的操作数，把每个转换后类型不同于源类型的操作数替换为按操作数生成的 builtin.unrealized_conversion_cast（标记 sc.unlowering），然后给整个源 op 标记 sc.unlowered 并将其重新插入为传递中合法。它不会擦除或替换该 op — 该 op 会存活下来，只是现在经过类型桥接。

// LowerToMlo EnqueueDMAOp lambda — @ 0x13239e00
// (lower_to_mlo.cc:405; same shape for EnqueueIndirectDMA / WaitDMA2 / FetchAndAddSync)
LogicalResult lower_enqueue_dma(EnqueueDMAOp op, EnqueueDMAOpAdaptor adaptor,
                                ConversionPatternRewriter& rw):
    srcOperands = op.getOperands()             // type-pointer list @ op+0x48
    newOperands = adaptor remapped operands    // ValueRange (already type-converted)

    // --- fast pre-scan: are ALL operand types already equal? ---
    //   compares (type_ptr & ~7) of src[i] vs remapped[i]  (line 102 / 137)
    if every src operand type == its remapped type:
        goto NO_REWRITE                         // nothing to bridge; op already legal

    rw.startRootUpdate(op)                       // [rw_vtable+0x28] (line 111)
    for i, (srcVal, newVal) in zip(srcOperands, newOperands):
        if (srcVal.type & ~7) != (newVal.type & ~7):          // type changed
            loc  = LocationGenerator::Visitor(op.getLoc())     // sc-location synthesis
            castTy = TypeRange{ newVal.type }                  // 1 result type (line 162)
            cast = UnrealizedConversionCastOp::create(rw, loc, castTy, srcVal)  // (line 175)
            cast->setAttr("sc.unlowering", UnitAttr::get(ctx))                  // (line 176-182)
            // splice: replace srcOperand[i] use with the cast result (intrusive list, line 183-199)
            op.setOperand(i, cast.getResult(0))
    op->setAttr("sc.unlowered", UnitAttr::get(ctx))            // mark source op (line 212-219)
    rw.finalizeRootUpdate(op)                    // [rw_vtable+0x30] (line 221)

NO_REWRITE:
    // trace-region tagging (independent of the bridge)
    if pass.insideTraceRegion:                   // v62[71] (line 224)
        op->setAttr("sc.inside_trace_region", UnitAttr::get(ctx))   // (line 226-232)
    return success                               // return 1 (line 234)
```text

### 三个属性标签

该发射最多放置三个 `UnitAttr` 标记 — 它们都已确认是由第二阶段 pass 消费的 `.rodata` 字符串字面量：

| 属性 | 设置在 | 含义 |
|---|---|---|
| `sc.unlowering` | 每个发射出的 `builtin.unrealized_conversion_cast` | “这个 cast 把尚未 lowering 的操作数桥接进 Mlo 类型系统；第二阶段 pass 拥有它” |
| `sc.unlowered` | 源 `tpu.enqueue_dma` op 自身 | “这个 op 被 LowerToMlo 有意留下未 lowering；在 tile expansion 期间展开它” |
| `sc.inside_trace_region` | 源 op，仅当 pass 位于 trace region 中时 | 传播 trace 嵌套，使展开后的 DMA 仍处于 trace 作用域内 |

> **注意 — bridge 在 `ExpandTiledMemRefsPass` 中解析，而不是 `LowerToSparseCoreLlvm`。** 反编译的 `0x13239e00` 函数体在固定位置设置标签：按操作数 cast 上的 `sc.unlowering`（`UnitAttr`，第 182 行）、源 op 上的 `sc.unlowered`（第 219 行），以及嵌套时的 `sc.inside_trace_region`（第 232 行）。下游消费者是由 `ExpandTiledMemRefsPass` 注册的 `substituteUnloweringConversionCastOp` (`0x134e73e0`) — 因此 bridge 的解析比标签名可能暗示的 `LowerToSparseCoreLlvmPass` (`0x13566d00`) 早一个阶段。
>
> **易错点 —** bridge 是**按操作数选择性**生成的，而不是整 op。预扫描（第 95-108 行）和内层循环（第 137 行）会比较每个源操作数与其 remapped 操作数的低 3 位屏蔽后类型指针；只有类型确实改变的操作数才会获得 cast。若一个 `enqueue_dma` 的所有操作数都保持类型不变，则**不会**发射 cast，只会获得 `sc.unlowered` 标记。无条件 cast 每个操作数的重新实现会发射死的 identity cast，第二阶段 pass 随后还必须剥离它们。
>
> **注意 —** 该 lambda 使用 `startRootUpdate` / `finalizeRootUpdate`（原地修改协议，`rw_vtable+0x28` / `+0x30`），而不是 `replaceOp`。因此该变更会被记录为 `ModifyOperationRewrite` 记录（[ConversionPatternRewriter](conversion-pattern-rewriter.md) 上的原地修改记录），插入的 cast 则记录为 `CreateOperationRewrite` + `UnresolvedMaterializationRewrite` 记录 — 所以整个 bridge 发射在 speculative legalizer 下是 rollback-safe 的。

---

## 阶段 1b — 块签名转换（`func::FuncOp` 合法性谓词）

### 目的

在到达任何 DMA op 之前，LowerToMlo 必须决定哪些函数需要让其*签名*经过 `TypeConverter` 重写。SparseCore 代码被组织为 sequencer 函数，其参数/结果类型（memref memory-space attr、`TupleType` 1:N 对、`WordType`）必须转换；普通函数则保持不变。这是 lowering 的块签名侧，并且完全由一个动态合法性谓词驱动。

### 三个动态合法性谓词

`runOnOperation` (`0x1322b200`) 通过 `addDynamicallyLegalOp` 安装三个合法性回调（调用普查显示四次 `addDynamicallyLegalOp` 调用 + 两次静态 `setOpAction`；这四次解析为下面三个不同谓词 lambda，其中一个注册给两个 op 类）。三者都返回一个 16 位值，其中 bit 0 = 当前合法，bit 8 (`0x100`) = “答案存在”（`ConversionTarget::isLegal` 读取的形式）。

| 谓词 | Lambda | 守卫的 Op |
|---|---|---|
| `$_0` FuncOp 签名 | `0x13231300` | `func::FuncOp` |
| `$_1` 传递中的 cast | `0x13231560` | `builtin.UnrealizedConversionCastOp` |
| `$_2` 结果类型 catch-all | `0x132315e0` | `OpResultTypeConversionPattern` 目标 |

#### `$_0` — `func::FuncOp` 合法性 (`0x13231300`)

```c
// func::FuncOp dynamic legality — @ 0x13231300
optional<bool> func_is_legal(FuncOp op):
    core   = TPUDialect::GetCoreTypeAttr(op)         // 0x14aa6020
    seqTy  = LowerMemrefToMlo::getSequencerType(op)  // 0x13507760
    // is this a SparseCore SEQUENCER function?  (op-name compares vs "sc"/"execute")
    if not is_sequencer(core, seqTy):
        return legal                                  // 0x100 — no signature rewrite
    // it IS a sequencer: walk the signature, every type must already be converted
    for ty in FunctionType.getInputs() ++ getResults():
        if typeConverter.convertType(ty) != ty:
            return illegal                            // signature still needs conversion
    return legal_recursive                            // 0x101 when input/result lists empty

一个函数除非是 SparseCore sequencer 函数，否则按原样合法（无需签名重写）— sequencer 通过 TPUDialect::GetCoreTypeAttr (0x14aa6020) 加 LowerMemrefToMlo::getSequencerType (0x13507760) 检测，二者会用 ASCII token "sc" 和 "execute" 测试函数的 core-type 属性和名称。对于 sequencer 函数，该谓词遍历 getInputs() ++ getResults() 并对每个类型调用 TypeConverter::convertType：函数合法当且仅当每个参数和结果类型都转换为自身。如果任何类型仍需转换，则 FuncOp 非法，这会触发 FuncOp 签名转换 pattern (0x13231ca0) 通过 converter 重写块签名。

注意 — FuncOp 签名转换 pattern (0x13231ca0) 执行实际的块参数重写。1:N 类型展开（一个 TupleType 参数拆成多个 SparseCore 标量、一个 I32Pair）会驱动块参数增殖；这些块级编辑的 rollback（BlockTypeConversionRewrite）是 ConversionPatternRewriter 上记录的动作日志的块重写部分。

`$_1` — `UnrealizedConversionCastOp` 合法性 (`0x13231560`)

这是允许 bridge-cast 在传递中穿过 legalizer 的谓词。没有它，DMA lambda 刚发射出的 unrealized cast 自身会被报告为非法，pass 会中止。

// UnrealizedConversionCastOp dynamic legality — @ 0x13231560
optional<bool> cast_is_legal(UnrealizedConversionCastOp op):
    attr = op.getInherentAttr("sc.unlowering")        // len 0xd; falls back to DictionaryAttr lookup
    if attr is absent:
        return 0x100                                  // "answer present", low bit clear → ILLEGAL
    // legal iff that attribute is the UnitAttr the bridge emits
    return (attr.getTypeID() == UnitAttr::id) | 0x100
```text

cast **只有在携带 `sc.unlowering` `UnitAttr` 时才合法** — 这正是 DMA bridge lambda 在 `0x13239e00` 给每个按操作数 cast 放置的标签。任何没有该标签的 unrealized cast（例如 type-converter 在签名边界插入的零散 reconciliation cast）都会被报告为非法（`0x100` 且低位清零），并且必须由 rewriter 物化。因此该谓词与阶段 1 发射形成直接的锁钥配对：lambda 写入 `sc.unlowering`，该谓词只接纳这些 cast。

#### `$_2` — `OpResultTypeConversionPattern` 目标合法性 (`0x132315e0`)

```c
// runOnOperation $_2 — @ 0x132315e0
optional<bool> result_types_legal(Operation* op):
    // fast accept: op's own dialect is in the target's legal-dialect DenseSet
    if op.getDialect().getTypeID() in MloConversionTarget.legalDialects:
        goto CHECK_RESULTS
    // the four DMA bridge ops are legal iff already tagged sc.unlowered (UnitAttr)
    if op is {EnqueueDMAOp, EnqueueIndirectDMAOp, WaitDMA2Op, FetchAndAddSyncOp}:
        u = op.getInherentAttr("sc.unlowered")        // len 0xc
        return u && (u.getTypeID() == UnitAttr::id) ? 0x101 : 0x100
    // memref-reshaping ops (EraseLayout / MemRefSqueeze / MemRefBitcast /
    //   MemRefSlice / ReinterpretCast / AssumeMultiple) are legal only when
    //   no operand/result still carries a tpu::TiledLayoutAttr memref
    ...                                               // TiledLayoutAttr guards
CHECK_RESULTS:
    for ty in op.getResultTypes():
        if typeConverter.convertType(ty) != ty:
            return 0x100                              // a result still needs conversion
    return 0x101                                      // legal (0x101 when no results)

catch-all OpResultTypeConversionPattern（MatchAnyOpTypeTag pattern，vtable 0x21903b40，matchAndRewrite 0x132456e0 → mlir::convertOpResultTypes 0x1c9572c0）是针对显式 pattern 未覆盖的任何 op 的通用 1:1 结果类型修复器：当某个 op 的操作数已被 remap 时，该 pattern 会转换该 op 的结果类型。$_2 是它的合法性门 — 在方言 / bridge-op / tiled-layout 快路径之后，一个 op 合法当且仅当 TypeConverter::convertType 把每个结果类型都映射为自身（即没有结果仍需转换）。

注意（低置信度）— $_2 不是纯结果类型谓词：0x132315e0 的反编译函数体首先基于 op 的方言是否属于 MloConversionTarget 合法方言集合进行短路，然后特判四个 DMA bridge op（一旦带有 sc.unlowered 标签即合法）和 memref-reshaping op（受 tpu::TiledLayoutAttr 守卫），只有 fall-through 路径才运行 convertType(ty) == ty 结果类型循环。上面的伪代码是已验证的形态；tiled-layout 守卫的精确分支顺序是摘要，并非逐字节解码。

阶段 2 — 延迟 DMA 物化解析

目的

bridge-cast 只是故事的一半；这个延迟最终必须被解析。本单元钉住带标签的 cast 在何处以及如何变成真正的 SparseCore DMA — 闭合 LowerToMlo lambda 打开的循环。

解析

applyFullConversion（LowerToMlo 驱动器）结束时，四个 DMA op 仍然存在，现在被标记为 sc.unlowered；它们发生改变的操作数由带 sc.unlowering 标签的 builtin.unrealized_conversion_cast op 供给；cast 合法性谓词 ($_1) 已把这些 cast 声明为传递中合法，因此该 pass 不会因它们而中止。

后续 SparseCore pass ExpandTiledMemRefsPass 拥有第二阶段。它为正好四个 bridge op 注册 conversion pattern：

第二阶段注册项	VA
`ExpandTiledMemRefsPass::addPattern<tpu::EnqueueDMAOp>`	`0x134eef60`
`ExpandTiledMemRefsPass::addPattern<tpu::EnqueueIndirectDMAOp>`	`0x134ef340`
`ExpandTiledMemRefsPass::addPattern<tpu::WaitDMA2Op>`	`0x134ef720`
`ExpandTiledMemRefsPass::addPattern<tpu::FetchAndAddSyncOp>`	`0x134f1240`
`sparse_core::expandTPUFetchAndAddSync(FetchAndAddSyncOp, adaptor, rw)`	`0x134e60c0`
`sparse_core::substituteUnloweringConversionCastOp(UnrealizedConversionCastOp, adaptor, rw)`	`0x134e73e0`

cast 解析器 substituteUnloweringConversionCastOp (0x134e73e0) 会找到带 sc.unlowering 标签的 cast 并替换它 — 一旦 tiled-memref 布局已知，就撤销这座桥，使现在具备形状的 DMA op 能展开为真正的 SparseCore 传输 op（rolled / retiled / granule-decomposed 传输链）。因为 ExpandTiledMemRefsPass 在 memref tiling 解析之后运行，所以这是 DMA 最终形态第一次可计算的时点 — 也正是 LowerToMlo 延迟它的原因。

注意 — 真正的 DMA 发射（LloRegionBuilder::EnqueueDmaGeneral / EnqueueDmaInGranules / EnqueueDmaToHostIova 族，0x1d543600 / 0x1d546700 / 0x1d548b20，以及 rolled/retiled 传输 helper）是第二阶段的内容，由 SparseCore-LLVM lowering 页面负责。本页的主张止于桥接契约：写入什么标签，哪个 pass 读取它，哪个解析器替换它。 granule 级传输代数只交叉引用，不重新推导。

端到端跟踪

text

  STAGE 1 — LowerToMlo (ModuleOp pass, applyFullConversion)
    tpu.enqueue_dma %src, %dst, %sflag, ...                 (operands in tpu types)
      └─ lambda 0x13239e00:
           per changed operand → %c = builtin.unrealized_conversion_cast %x
                                       {sc.unlowering}        (→ Mlo type)
           op.setAttr("sc.unlowered")                        (+ sc.inside_trace_region?)
      ⇒ op SURVIVES, type-bridged; cast declared legal by $_1 (0x13231560)
  ── applyFullConversion succeeds with the bridge intact ──
  STAGE 2 — ExpandTiledMemRefsPass (after memref tiling resolved)
    addPattern<EnqueueDMAOp> 0x134eef60
      └─ substituteUnloweringConversionCastOp 0x134e73e0   (resolve sc.unlowering casts)
      └─ expand sc.unlowered DMA → real SparseCore transfer ops
           (EnqueueDmaGeneral / InGranules / ToHostIova: 0x1d543600 / 0x1d546700 / 0x1d548b20)
```text

LowerToMlo DMA Bridge-Cast ​

摘要 ​

为什么分两阶段 — 直接 lowering 与桥接的分割 ​

目的 ​

分割 ​

阶段 1 — DMA Bridge-Cast（tpu.enqueue_dma → 带标签 cast） ​

目的 ​

算法 ​

$_1 — UnrealizedConversionCastOp 合法性 (0x13231560) ​

阶段 2 — 延迟 DMA 物化解析 ​

目的 ​

解析 ​

端到端跟踪 ​

LowerToMlo DMA Bridge-Cast

摘要

为什么分两阶段 — 直接 lowering 与桥接的分割

目的

分割

阶段 1 — DMA Bridge-Cast（`tpu.enqueue_dma` → 带标签 cast）

目的

算法

`$_1` — `UnrealizedConversionCastOp` 合法性 (`0x13231560`)

阶段 2 — 延迟 DMA 物化解析

目的

解析

端到端跟踪