Mosaic 概览

本页中的所有地址、符号、op 名称字符串、dump 阶段名和错误字符串均适用于来自 libtpu-0.0.40-cp314 wheel 的 libtpu.so（build-id 89edbbe81c5b328a958fe628a9f2207d，构建 libtpu_lts_20260413_b_RC00）。其他版本会有所不同。

摘要

Mosaic 是通过一个 HLO opcode 进入的平铺张量 MLIR kernel 编译器，即 kCustomCall("tpu_custom_call")，并且它是 TPU 设备路径上 mlir::tpu dialect 的唯一生产者。Pallas @pl.kernel（或手写的 Mosaic kernel）会在 libtpu 之外降低为一个 tpu-dialect MLIR module；该 module 被序列化、嵌入 custom-call 的 backend config，并随 HLO 程序跨过 PjRt 边界传递。当编译器在 HLO→LLO 发射阶段到达该 custom-call 时，libtpu 会解析嵌入的 module，运行一次版本迁移往返，并驱动它经过 16 阶段的 tpu-dialect pass pipeline 降到 LLO（或者对于 SparseCore kernel，降到 LLVM-TPU intrinsics）。普通 HLO 永远不会变成 tpu ops，而是由约 3225 个 xla::jellyfish::*Emitter 类直接发射到 LLO（MHLO → XTile → tpu）。因此，Mosaic 是一个向同一目标 dialect 供给 IR 的第二生产者，而不是并行的后端：它跳过整个 HLO optimizer，并在 tpu dialect 处重新加入主下降路径（The tpu MLIR Dialect）。

这是 Mosaic 后端的导览页。它明确 Mosaic 是什么、将 tpu module 带入编译器的 import/serde 接缝、把导入 module 粘接到周围 HLO 程序的 CustomCallEmitter::Emit driver、16 阶段 RunMLIRPasses pipeline（TensorCore 与 SparseCore 两条分支），以及 Mosaic 子页面地图。它不会复述 layout 代数，VectorLayout (sublane, lane) 值类型位于 Mosaic VectorLayout，逐 op 的 layout 推断规则位于 Mosaic Layout Inference；这里只链接，不重复。

对于重新实现，导览契约是：

Mosaic 是唯一的 tpu-dialect 来源。 二进制中有零个 *ToTpuPass/MhloToTpu/StablehloToTpu 转换 pass；tpu dialect 是被导入的，由上游创作，绝不是从 MHLO 降低而来。重新实现者不应构建并不存在的 MHLO→tpu legalizer。
import + serde 接缝。 GetMlirModuleOpFromCustomCall → GetCachedCustomCallBody（以 CityHash128 为 key，每个唯一 kernel 只解析一次）→ GetMlirModule → ParseModule + 单次 MosaicSerdePass deserialize/upgrade。serde 是逐 op 的版本迁移引擎，不是通用 codec。
driver 入口。 CustomCallEmitter::Emit 会依据 HLO operands 验证 kernel main 的 ABI，绑定 memory-space colors 与 windows，按 core type 选择入口 func，运行 pipeline，并把生成的 LLO 缝合到父 region。
16 阶段 pipeline。 RunMLIRPasses 运行 simplify → infer-memref-layout → canonicalize-mosaic → infer-vector-layout → apply-vector-layout → … → lower-to-llo，其中 SparseCore 分支会转向 lower-to-mlo → LLVM-TPU。
子页面地图。 layout 代数、layout 推断规则、tpu→LLO 下降，以及 SparseCore lowering 各自所在的位置。


Mosaic 是什么	平铺张量 MLIR kernel 编译器；TPU 路径上 `mlir::tpu` 的唯一生产者，通过 `tpu_custom_call` 导入
入口 opcode / 目标	`HloOpcode::kCustomCall`, `custom_call_target = "tpu_custom_call"`
Driver 入口	`xla::jellyfish::(anon)::CustomCallEmitter::Emit` @ `0x111ef740`（3657 行；源码 `…/mosaic/python/custom_call_emitter.cc`）
Import 接缝	`mlir_utils::GetMlirModuleOpFromCustomCall` @ `0x13e327a0` → `GetCachedCustomCallBody` @ `0x13e31860` → `GetMlirModule` @ `0x13e31220`
Serde 引擎	`mlir::tpu::MosaicSerdePass::runOnOperation` @ `0x145307a0`；`jaxlib::mosaic::RunSerde` @ `0x14533b20`；module attr `stable_mosaic.version`，默认版本 11
Pipeline driver	`xla::jellyfish::RunMLIRPasses` @ `0x111fefa0`（源码 `…/mosaic/python/mosaic_passes.cc`）；16 个阶段，dump `post-<stage>`
入口 func 选择	`mlir::tpu::GetFuncWithCoreType` @ `0x14aa61a0`（匹配 `tpu.core_type`，回退到 `@main`）
Mem-space 桥接	`MemorySpaceToColor` @ `0x1d6ffb80` / `ColorToMemorySpace` @ `0x1d6ffb00`
TensorCore 交接	`mlir::tpu::createLowerToLLOPass` @ `0x11203ba0` → `llo::createEliminateLLOExtensionsPass` @ `0x13e668a0`
SparseCore 交接	`mlir::tpu::createLowerToMloPass` @ `0x1322adc0` → `sparse_core::CreateLowerToSparseCoreLlvmPass` @ `0x135667c0`
二进制中的 MHLO→`tpu` passes	无（0 个 `*ToTpuPass`/`MhloToTpu`/`StablehloToTpu` 符号）
置信度	HIGH（由符号/字符串锚定），除非某行或标注另有说明

Mosaic 是什么，以及为什么它是一条侧通道

Mosaic 不是 HLO optimizer 内部的一个阶段；它是一个逃逸通道。通用 TPU 计算路径通过约 3225 个 xla::jellyfish::*Emitter 类（DmaEmitter、FusionEmitter、ReduceEmitter、GatherEmitter 等）调用 LloRegionBuilder（3037 个引用函数），直接把 HLO 降低到 LLO；对于通用计算，永远不会产生 MLIR tpu dialect。二进制中存在 tpu dialect 只是因为 Pallas/Mosaic kernel 在上游用它创作并被导入。两个独立观察可以钉住这一点：

没有转换 pass 产生 tpu。 函数表中匹配 *ToTpuPass、*ToTPUDialect、MhloToTpu、HloToTpu、StablehloToTpu、LegalizeToTpu、LowerMhloToTpu 的符号为零。（仅有的 ConvertToTpu* 命中是 ConvertToTpuCompilationCacheGroupEntryProto 和 ConvertToTpuCoreProgram，属于 runtime/SparseCore-ISA，无关。）
通用路径完全绕过 tpu。 每个 *Emitter 都是逐 HLO opcode 的 visitor，直接在 LloRegion 中构建 llo.* ops。

因此，设备路径是两棵树，只有 Mosaic 分支在 tpu dialect 处汇合：

text

TREE A1 — general TPU compute
  HLO ──[97-row pre-passes]──► HLO ──[~3225 jellyfish *Emitter + LloRegionBuilder]──► LLO ──► bundles
        (never becomes tpu-dialect ops)

TREE A2 — Pallas/Mosaic kernels  ◄── THIS PAGE
  HLO kCustomCall("tpu_custom_call") carrying a SERIALIZED tpu-dialect module
        ──[import + serde]──► tpu-dialect IR
        ──[RunMLIRPasses 16-stage pipeline]──► LLO (TensorCore) | LLVM-TPU (SparseCore)
```text

Mosaic kernel 跳过 HLO optimizer（kernel body 没有 algebraic simplification、layout assignment、fusion 或 MSA），因为它们到达时*已经*是带有自身 tiling 的 `tpu` ops。它们在 `tpu` dialect 处与优化后的路径汇合，并从那里共享相同的 tpu→LLO→bundle 下降。

> **GOTCHA — 不要构建 MHLO→`tpu` legalizer。** [compiler overview](overview.md) 在 Level 2（`tpu` dialect）画出的汇合点只有一个生产者供给：Mosaic custom-call import。普通 HLO 绕过 `tpu` 并直接进入 LLO。把精力投入 MHLO/StableHLO→`tpu` conversion-pattern 集合的重新实现，是在构建生产编译器中不存在的阶段。完整的两树证据见 [MHLO → XTile → tpu](mhlo-xtile-tpu-lowering.md)。
>
> **NOTE — "TLP"/"MLO" 是 IR 层名称，不是 dialect。** `tpu` dialect 作为 "TPU-Level Program" 容器的物理角色只对导入的 Mosaic kernel 实现。"Mlo"（mid-level ops）是 `tpu`/`sparse_core` lowering **内部**的 SparseCore 中间层（见 [The SparseCore Arm](#the-sparsecore-arm)），不是顶层 IR。

---

## Mosaic Kernel 嵌入

Pallas `@pl.kernel` 由 JAX frontend（在 `libtpu` 之外）编译成一个 `tpu`-dialect module，以**serialize** 模式运行 `MosaicSerdePass`（降级到固定版本，设置 module 的 `stable_mosaic.version` attr），并嵌入到一条 HLO instruction：

```text
HloOpcode::kCustomCall
  custom_call_target = "tpu_custom_call"
  backend_config     = serialized xla::jellyfish::CustomCallConfig proto:
      mlir_module           (an absl::Cord; the serialized tpu module)
      serialization_format  (must equal 1)
      input_memory_colors / output_memory_colors
      has_communication, collective_id, cost_estimate, metadata

四个 HLO pre-pass 会在 lowering 前准备该调用（Custom-Call Lowering）：MosaicFusion（把周围 HLO 融合进 operands）→ TpuCustomCallLegalizer（TensorCore/SparseCore/Megachip 分类）→ TpuCustomCallMemorySpacePolicy（填充 memory colors）→ TpuCustomCallScopedVmemAdjuster（调整 scoped VMEM 大小）。在 HLO→LLO 发射时，"tpu_custom_call" 目标注册的 emitter 是 CustomCallEmitter::Emit。

Import 与 MosaicSerde 往返

import 是一个很薄的缓存包装器，它对每个唯一 kernel 只解析一次嵌入的 module，并运行一次 deserialize/upgrade pass。

text

GetMlirModuleOpFromCustomCall (0x13e327a0)
  └─ GetCachedCustomCallBody (0x13e31860)
        1. validate body present (else "Custom call body is empty.")
        2. key = farmhash::CityHash128WithSeed(Cord→string(mlir_module))   // 128-bit
        3. HloModule::GetCacheEntry<MosaicMlirCacheEntry>(module, key)
              HIT  → reuse the parsed ModuleOp across every duplicate kernel
              MISS → build it:
                       ctx = new mlir::MLIRContext
                       LoadDialects(ctx)                      // 0x13e32140
                       GetMlirModule(config, ctx, verify)     // 0x13e31220
                       cache a MosaicMlirCacheEntry (OWNS ctx + ModuleOp)
        4. fingerprint-collision guard → fatal on bcmp mismatch
              ("Kernel body fingerprint collision detected for key: …")
```text

`LoadDialects`（`0x13e32140`）构建 kernel 的*自包含* MLIRContext（它不共享 host compiler 的 context）：加载 `stablehlo`、`arith`、`func`、`memref`、`scf`、`vector`、**`mlir::tpu::TPUDialect`**、`cf`、**`mlir::llo::LLODialect`**、`mlir::sparse_core::ScDialect`（注册为 `"sc_tpu"`），以及 `RegisterNonSparseCoreDialects` 和 `func::registerAllExtensions`，随后调用 `allowUnregisteredDialects(true)`；serde 的 serialize 侧需要它，也需要它来让未知 ops 在往返中存活。

`GetMlirModule`（`0x13e31220`）解析并反序列化：

1. `ParseModule`（`0x13e30dc0`）包装 `mlir::parseSourceString`，通过同一入口消费**文本 MLIR 和 MLIR bytecode**，位于 `StatusScopedDiagnosticHandler` 下，然后执行 `constructContainerOpForParserIfNecessary<ModuleOp>`。失败：`"Failed to parse the Mosaic module: <diag>"`。
2. 如果 `serialization_format == 1`，在 `PassManager` 中运行一次 `MosaicSerdePass(serialize=false)`（dump 阶段 `"deserialization"`）；任何其他值 → `"Unsupported serialization format: <N>"`。

### MosaicSerdePass 是版本迁移引擎

`MosaicSerdePass::runOnOperation`（`0x145307a0`）**不是**通用 bytecode codec，而是逐 op 的版本升级/降级引擎。它以两种模式运行：`serialize=false`（deserialize/**upgrade**，导入时使用）和 `serialize=true`（serialize/**downgrade**，把 kernel 写入持久 compilation cache 时使用）。它惰性构建两个按 op-name 作为 key 的 `StringMap`（`upgrade_rules`、`downgrade_rules`），由 `jaxlib::mosaic::RunSerde`（`0x14533b20`）分派。

`RunSerde` 读取 module 的 `stable_mosaic.version` IntegerAttr，检查边界（`"Unsupported version: expected <= <cur> but got <stored>"`），移除 attr，然后遍历每个 op 并应用匹配的迁移 lambda。八个 op 带有逐版本规则：

| op (key) | 该规则编码的 schema 演化 |
|---|---|
| `tpu.enqueue_dma` | v2 增加了 operand-segment-sizes 形式；v4 增加 `priority`；v8 增加 `strict_ordering` |
| `tpu.wait_dma2` | v7 增加 `device_id`/`core_id` segments |
| `tpu.dynamic_gather` | v5：单数 `dimension` → 多维 `dimensions` |
| `tpu.iota` | v6：单数 `dimension` → 多维 `dimensions` |
| `tpu.sem_signal` | v2 增加 `core_id` |
| `vector.multi_reduction` | v3：`reduction_dims` ArrayAttr-of-IntegerAttr → DenseArrayAttr&lt;i64&gt; |
| `tpu.store` | v11 增加 fused store-and-accumulate 的 `add` flag |
| `arith.constant` | v10 泛化 vector-of-i1 constants（v10 之前必须是 splat） |

当前 dialect 版本是 **11**。读取 JAX 发出的持久 cache 的 out-of-tree compiler 必须向下重放匹配的 downgrade lambdas，直到自己的固定版本。字符串常量解析为 `kMangledDialect = "stable_mosaic"`（bytecode dialect tag，13 个字符）和 `kVersionAttrName = "stable_mosaic.version"`。逐版本的 field-migration body 与 layout serialization 一起记录在 [Mosaic VectorLayout](mosaic-vectorlayout.md)。

> **NOTE — bytecode magic 属于上游 MLIR；TPU 专属部分是 version attr + 8 条规则。** 解析处理标准 MLIR text/bytecode；唯一 TPU 专属的反序列化工作是读取/移除 `stable_mosaic.version` 并重放八个逐 op upgrade lambda。

---

## Driver：CustomCallEmitter::Emit

`CustomCallEmitter::Emit`（`0x111ef740`，3657 行）是 `"tpu_custom_call"` 注册的 LLO emitter。它把导入的 `tpu` module 连接到周围 HLO 程序，并驱动整个 pipeline。恢复出的调用结构和逐字错误字符串：

1. **获取并验证 config** — `"Custom call does not have a custom call config."` / `"Failed to retrieve the config. Error: %s"`。
2. **导入 body** — `GetCachedCustomCallBody` → 缓存的 `tpu` `ModuleOp`。重入保护：`"Trying to lower the same Mosaic kernel more than once"`。
3. **Core-type gate**（TensorCore vs SparseCore）来自 cache entry 的 core-type 集合：`"Cannot lower a Mosaic kernel targeting SparseCore on a device without SparseCore"`、`"Cannot lower a Mosaic kernel targeting SparseCore to TensorCore."`、`"Mosaic is not supported on this hardware version."`。
4. **验证 kernel `main` ABI**（HLO↔kernel 契约）：

   ```text
   main.getNumArguments() == iteration_bounds + num_inputs + num_outputs
                             + num_scratch + has_communication
   main.getNumResults()   == 0          (all outputs are by-ref memrefs)
   main.getRegion().hasOneBlock()

绑定 operands/outputs — element-type allow-list（float8e5m2、float8e4m3fn、float8e4m3b11fnuz、float8e8m0fnu、float4e2m1fn、bfloat16、32-bit）、逐 operand memory-space（"Some arguments are missing their memory space attribute"）、通过 MemorySpaceToColor 的 memory-space→color、通过 VerifyAndMaybeUpdateShapeLayoutForMosaic 的 layout 验证（"Missing layout for Mosaic kernel operand %d" / "Failed to verify layout for Mosaic kernel operand "）、scratch placement（"Scratch memref allocation only supported for vmem, smem and semaphore_mem."），以及退出时 semaphore 活性（"Semaphore (scratch argument %d) has a nonzero value upon exit … Make sure every DMA is awaited …"）。
Windowing — PipelineEmitter（在 Emit 中 39 个引用）把每个 operand/output 划窗到 iteration grid 中，从 main 读取 window_params、iteration_bounds 和 dimension_semantics attrs（"Wrong number of window_params: expected %d, got %d"、"iteration_bounds should be a DenseI64ArrayAttr"）。window-descriptor 数学（grid flatten、double-buffering、dynamic-offset DMA、window cycle cost）在 Mosaic VectorLayout 中恢复。
选择入口 func — GetFuncWithCoreType(module, core_type)（0x14aa61a0）返回 tpu.core_type 匹配设备的 func.func，回退到 @main（"No function with tpu.core_type = %v nor a main function found"）。
运行 pipeline — RunMLIRPasses(...) 就地把 module 的 main 降到 LLO（下一节）。
Megacore split — MegacoreAdjuster（4 个引用）把工作拆分到 megacore chip 的两个 TensorCore 上。
缝合进 LLO — 降低后的 body 通过 LloRegionBuilder（AllocateScopedVmem/Smem、EnqueueDmaLocalInGranules 等）发射进父 LloRegion；post-emit cleanup "post-finalize-llo-post-emitter"。返回 kernel outputs 的 LLO SSA value(s)。

NOTE — driver 入口是 CustomCallEmitter::Emit，不是 MosaicEmitter。 MosaicEmitter（EmitWindow @ 0xfaadcc0）是 anonymous-namespace 中一个很小的窗口发射 helper，是 MosaicBroadcast 和 PipelineEmitter::OperandWindow 的同级，在 windowing 步骤内部使用。真正导入 module、验证 ABI、运行 pipeline 并缝合 LLO 的类是 CustomCallEmitter::Emit（0x111ef740）。"MosaicEmitter" 是 Mosaic 侧门的宽泛叫法；driver 符号是 CustomCallEmitter::Emit。

导入 module 如何连接到 HLO

连接方式是 main 上的 ABI 加上 memory-space-color 桥：

周围 HLO 概念	`tpu`-module `main` 绑定
`kCustomCall` operand i（HLO buffer）	一个 `memref` arg（tiled layout + `tpu.memory_space` attr）；color 来自 `config.input_memory_colors[i]`
`kCustomCall` result	一个 by-ref output `memref` arg（`num_results == 0`）
grid / pipeline dims	`main` 上的 `iteration_bounds`（`DenseI64ArrayAttr`）
parallel vs reduction grid dims	`dimension_semantics` attr
逐 operand window/tiling	`window_params` attr（每个 memref arg 一个）
scoped scratch (vmem/smem/sem)	outputs 之后的额外 `main` args
collective barrier	`has_communication` arg（+ `config.collective_id`）

MemorySpaceToColor（0x1d6ffb80）把 mosaic-tpu memref memory-space enum 映射到 host BufferAssignment 使用的整数 "color"，从而让导入 kernel 的 inputs/scratch 落到正确的物理内存中。mosaic-tpu MemorySpace enum（vmem/smem/hbm/cmem/semaphore_mem/vmem_shared/host）和 jellyfish color table 值记录在 Mosaic VectorLayout。

tpu-Dialect Pass Pipeline：RunMLIRPasses

RunMLIRPasses（0x111fefa0，源码 …/mosaic/python/mosaic_passes.cc）是真正的 tpu-dialect pipeline。每个 pass 都通过 RunPass（0x14514d60）运行，后者用 RunAndCaptureDiagnostics 包装 PassManager，并将 MLIR dump 为 post-<stage>（或 post-<stage>-failed）。TensorCore pipeline 按执行顺序和 create-functions（除注明外均为 mlir::tpu::）如下：

#	stage (dump name)	create-function / action
0	`original`	`DumpMlir`；setup：`GetFuncWithCoreType`、`GetMosaicHardwareGeneration(target)`；guard `"Should not run HLO passes without layout passes"`
1	`deserialize`/diagnostics	`RunAndCaptureDiagnostics`
2	`simplify`	`$_4` canonicalizer → `InferMemorySpaces`（dump `post-infer-memref-space`）
3	`hlo-conversion` (iff `run_hlo_passes`)	嵌套 `func.func`：`stablehlo::createStablehloLegalizeToLinalgPass` + `LinalgVectorizationPass`
4	`infer-memref-layout`	`createInferMemRefLayoutPass`（`0x132c0f00`）（+ `-simplify`）
5	`assert-insertion`	`createDebugAssertInsertionPass`
6	`canonicalize-operations`	`createCanonicalizeOperationsPass`
7	`pre-canon-optimization`	`createPreCanonicalizationOptimizationPass`
8	`canonicalize-mosaic`	`createCanonicalizeMosaicPass`（`0x132a2ac0`）（+ `-simplify`）
9	`tiling-propagation`	`createTilingPropagationPass`
10	`infer-vector-layout`	`createInferVectorLayoutPass`（`0x132c2c20`）
11	`relayout-insertion`	`createRelayoutInsertionPass`
12	`apply-vector-layout`	`createApplyVectorLayoutPass`（`0x1325cda0`）（+ `-simplify`）
13	`logical-to-physical-device-id`	`createLogicalToPhysicalDeviceIdPass`
14	`lower-to-llo`	嵌套 `func.func`：`createLowerToLLOPass(target)`（`0x11203ba0`）
15	`eliminate-llo-extensions`	嵌套 `func.func`：`llo::createEliminateLLOExtensionsPass`（`0x13e668a0`）
16	`finalize-llo`	`$_4` canonicalizer

finalize-llo 之后，module body 是纯 llo.*（加上 scf/func/memref structural ops），控制返回 CustomCallEmitter::Emit 进行 LloRegionBuilder 缝合。-simplify 子阶段会在三个 layout pass 之后运行 $_4 canonicalize/CSE lambda。IR printing 受 VLOG 控制；verifier 受 IsMosaicVerificationEnabled 控制。

pipeline 的 layout 核心是阶段 10–12。infer-vector-layout 用 in_layout/out_layout attributes 标注每个 op（逐 op 推断规则见 Mosaic Layout Inference）；relayout-insertion 在 consumer 需要的 layout 与 producer 的 layout 不同时插入显式 tpu.relayout，使 apply-vector-layout（applyLayoutFunc @ 0x1325cc80 → 逐 op applyLayoutOp）能通过 49 项 op-name→rule 分发表，把每个逻辑 vector<NxMxT> 物化为硬件原生 vreg tiles。VectorLayout 值类型、文本语法、逐 vreg packing 数学和 relayout driver 都在 Mosaic VectorLayout。

GOTCHA — canonicalize-mosaic 以 allow_unverified=!a12 运行。 createCanonicalizeMosaicPass 接收 gen + IsMosaicCompatibilityModeEnabled + 一个从 verification gate 派生的 allow_unverified flag。重新实现者不应假设 canonicalizer 总是重新验证；在 compatibility mode 下，它会容忍未验证 IR，以便旧 kernel 存活。
NOTE — 已盘点的 22 个 tpu passes 中有 4 个不在 TensorCore 分支中。 CanonicalizeMemorySpace 和 ConvertIntegerMemrefs 属于其他路径，LowerToMlo 属于下面的 SparseCore 分支，MosaicSymbolicShapeRefinement 作为 hybrid HLO+MLIR pass 更早运行。在这份 16 阶段列表里搜索它们是找不到的。

TensorCore vs SparseCore：两条 Lowering 分支

pipeline 在 logical-to-physical-device-id 之后依据入口 func 的 tpu.core_type 分叉。

TensorCore 分支（阶段 14–16）

createLowerToLLOPass(target)（0x11203ba0）运行 tpu.* → llo.* 的完整 applyFullConversion（242 个 patterns，见 tpu → LLO Lowering）；createEliminateLLOExtensionsPass（0x13e668a0）把 LLO "extension" convenience-macro ops 展开成基础 LLO ops；finalize-llo 执行 canonicalize。此后，每个 vector value 都是原生 vreg shape，1:1 映射到 LLO VregType，body 被缝合进父 LloRegion。下游，TpuProgram Serialization 和 scheduling back end 会打包 bundles。

SparseCore 分支 {#the-sparsecore-arm}

main 带有 SparseCore tpu.core_type 的 kernel 不会降低到 LLO。相反，tpu ops 会转换到 SparseCore "Mlo"（mid-level-ops）阶段，然后转换到 LLVM-TPU intrinsics：

text

tpu (SparseCore core_type) module
  → [mosaic_sc::InferVectorLayoutPass + ApplyVectorLayoutPass]      (vreg layout; only op: mosaic_sc.relayout)
  → mlir::tpu::createLowerToMloPass(Target, LowerToMloPassContext)  (0x1322adc0)
        converts tpu.* (incl. SparseCore-only tpu.barrier, tpu.all_reduce,
        tpu.enqueue_indirect_dma, tpu.fetch_and_add_sync, tpu.scan,
        tpu.wait_indirect_dma) into mlo/sparse_core/llvm ops
  → MloModuleVerifier
  → sparse_core::CreateLowerToSparseCoreLlvmPass(Target, …)         (0x135667c0)  → mlir::LLVM / LlvmTpu
  → sparse_core::CreateLlvmIntToPtrSafetyPass(Target)               (int↔ptr safety legalize)
  → mlir-translate -mlir-to-llvmir → per-gen SparseCore ISA
```text

SparseCore Mosaic dialect（`mosaic_sc`）刻意保持很小，只有一个 op（`mosaic_sc.relayout`）加两个 layout passes，因此 Mosaic vreg-layout 阶段能在 `tpu → mlo → LlvmTpu` lowering 前为 SparseCore 运行。细节见：[LowerToMlo DMA Bridge](lower-to-mlo-dma-bridge.md)、[LowerToSparseCoreLlvm](lower-to-sparsecore-llvm.md)、[SCTypeConverter](sc-type-converter.md) 和 [SparseCore Part](../sparsecore/overview.md)。

---

## 完整的 Mosaic → tpu → LLO 链

```text
JAX/Pallas frontend (out of libtpu)
  @pl.kernel → tpu-dialect module → MosaicSerde SERIALIZE (downgrade, set stable_mosaic.version)
             → embedded in HloCustomCall("tpu_custom_call").backend_config.mlir_module
─────────────────────────────── libtpu.so ───────────────────────────────
HLO pre-passes: MosaicFusion → TpuCustomCallLegalizer
                → TpuCustomCallMemorySpacePolicy → TpuCustomCallScopedVmemAdjuster
   │  at HLO→LLO emit time, target = "tpu_custom_call"
   ▼
CustomCallEmitter::Emit (0x111ef740)
   │  GetMlirModuleOpFromCustomCall (0x13e327a0)
   │    └ GetCachedCustomCallBody (0x13e31860): CityHash128 key; on miss
   │        new MLIRContext → LoadDialects → GetMlirModule (parse + MosaicSerde deserialize)
   │  validate main ABI; bind memref colors (MemorySpaceToColor); set windows (PipelineEmitter)
   │  GetFuncWithCoreType (0x14aa61a0) → pick entry func
   │        ▼  RunMLIRPasses (0x111fefa0)
   │   simplify → infer-memref-layout → canonicalize-mosaic → tiling-propagation
   │   → infer-vector-layout → relayout-insertion → apply-vector-layout
   │   → logical-to-physical-device-id
   │   ├── TensorCore: → lower-to-llo (242 patterns) → eliminate-llo-extensions → finalize-llo
   │   └── SparseCore: → (mosaic_sc layout) → lower-to-mlo → MloModuleVerifier
   │                   → lower-to-sparsecore-llvm → llvm-int-to-ptr-safety ⇒ LLVM-TPU
   │  MegacoreAdjuster (split across 2 TensorCores)
   │  stitch LLO into parent LloRegion via LloRegionBuilder
   ▼
LLO instruction stream → bundle packer → ISA bundles   (SparseCore: LLVM IR → per-gen SC ISA)

Mosaic 子页面地图

topic	page	owns
layout 原子	Mosaic VectorLayout	`(sublane, lane)` `VectorLayout` struct + invariants、文本语法、`tilesPerVreg`/`tileArrayShape`、`ImplicitDim` 模型、`relayout` driver、49 项 `applyLayoutOp` 分发表、8 个 serde migration lambdas、`MemorySpaceToColor` 表值，以及 `PipelineEmitter` window 数学
layout 求解器	Mosaic Layout Inference	逐 op 的 `VectorLayoutInferer` 规则，用来选择每个 op 的 `in_layout`/`out_layout`（也就是 `applyLayoutOp` 消费的 attrs 的生产者）
目标 dialect	The tpu MLIR Dialect	`tpu` dialect 的 ops、attributes 和 op-model contract
TensorCore 交接	tpu → LLO Lowering	`createLowerToLLOPass` 到 LLO 的 242-pattern 下降
SparseCore 交接	LowerToMlo DMA Bridge、LowerToSparseCoreLlvm、SCTypeConverter	`tpu → mlo → LlvmTpu` SparseCore lowering
dispatch 层	Custom-Call Lowering	`tpu_custom_call` 目标注册 + 四个 HLO pre-passes

置信度总结

Claim	Evidence
Mosaic（`tpu_custom_call` import）是唯一的 `tpu`-dialect 生产者	0 个 `ToTpuPass`/`MhloToTpu`/`StablehloToTpu` 符号；通用路径有约 3225 个 jellyfish `Emitter` + 3037 个 `LloRegionBuilder` refs
Import 接缝 = `GetMlirModuleOpFromCustomCall` → `GetCachedCustomCallBody` → `GetMlirModule`/`ParseModule`/`LoadDialects`	`0x13e327a0`、`0x13e31860`、`0x13e31220`、`0x13e30dc0`、`0x13e32140` 已反编译；CityHash128 key + `MosaicMlirCacheEntry` cache
Serde 是逐 op 版本迁移引擎；8 个 ops；默认 v=11；attr `stable_mosaic.version`	`MosaicSerdePass::runOnOperation` `0x145307a0`，`RunSerde` `0x14533b20`；`kMangledDialect`/`kVersionAttrName` 来自 `.data` relocs
Driver 入口是 `CustomCallEmitter::Emit`（3657 行），调用 `GetCachedCustomCallBody`/`GetFuncWithCoreType`/`MemorySpaceToColor`/`PipelineEmitter`/`MegacoreAdjuster`/`RunMLIRPasses`	`0x111ef740` 已反编译；1/1/2/39/4/1 个引用已确认；逐字 ABI/memory/window/semaphore 错误字符串
`MosaicEmitter` 是 window-emit helper，不是 driver	`MosaicEmitter::EmitWindow` `0xfaadcc0`，是 `MosaicBroadcast`/`PipelineEmitter::OperandWindow` 的同级；不是 lowering driver
16 阶段 `RunMLIRPasses` pipeline，带列出的 stage names + create-functions	`0x111fefa0`（源码 `mosaic_passes.cc`）；`createInferMemRefLayoutPass` `0x132c0f00`、`createCanonicalizeMosaicPass` `0x132a2ac0`、`createInferVectorLayoutPass` `0x132c2c20`、`createApplyVectorLayoutPass` `0x1325cda0`、`createLowerToLLOPass` `0x11203ba0`、`createEliminateLLOExtensionsPass` `0x13e668a0`
入口 func 通过 `GetFuncWithCoreType` 按 `tpu.core_type` 选择	`0x14aa61a0`；错误字符串 `"No function with tpu.core_type = %v …"`
SparseCore 分支 = `mosaic_sc` layout → `createLowerToMloPass` → `MloModuleVerifier` → `CreateLowerToSparseCoreLlvmPass` → LLVM-TPU	`createLowerToMloPass` `0x1322adc0`，`CreateLowerToSparseCoreLlvmPass` `0x135667c0`，`mosaic_sc::createInferVectorLayoutPass` `0x132ecf60`
每个 TPU generation 上精确的 `hlo-conversion`/megacore/large-2nd-minor 分支 gating	调用序列已恢复；逐 generation flag 选择未完全追踪

交叉引用

The TPU Compiler (overview) — 五阶段主干；本页详述其 "Mosaic side channel" 部分。
MHLO → XTile → tpu Lowering — 两树证据，以及不存在 MHLO→tpu pass 的证明；通用路径的 jellyfish *Emitter 路线。
Custom-Call Lowering & the Target Registry — tpu_custom_call 目标注册，以及准备嵌入的四个 HLO pre-passes。
The tpu MLIR Dialect — Mosaic 导入且 pipeline lowering 的目标 dialect。
Mosaic VectorLayout — (sublane, lane) layout 代数、applyLayoutOp 规则表、serde migration lambdas，以及 window-descriptor 数学（这里只链接，不重复）。
Mosaic Layout Inference — 产生 in_layout/out_layout attrs 的逐 op 推断规则。
tpu → LLO Lowering — TensorCore createLowerToLLOPass 下降。
LowerToMlo DMA Bridge · LowerToSparseCoreLlvm · SCTypeConverter — SparseCore lowering 分支。
TpuProgram Serialization — 降低后的 LLO 如何被打包和序列化。
Part VIII — Scheduling · SparseCore Overview — 两条分支交接到的后端。
Binary: extracted/libtpu-0.0.40-cp314-cp314-manylinux_2_31_x86_64/libtpu/libtpu.so（build-id 89edbbe81c5b328a958fe628a9f2207d）
Index entry: Part V — Compiler: Lowering & Optimization Passes / MLIR lowering chain — back to index

Mosaic 概览 ​

摘要 ​

Mosaic 是什么，以及为什么它是一条侧通道 ​

Import 与 MosaicSerde 往返 ​

导入 module 如何连接到 HLO ​

tpu-Dialect Pass Pipeline：RunMLIRPasses ​

TensorCore vs SparseCore：两条 Lowering 分支 ​

TensorCore 分支（阶段 14–16） ​

SparseCore 分支 {#the-sparsecore-arm} ​

Mosaic 子页面地图 ​

置信度总结 ​

交叉引用 ​

Mosaic 概览

摘要

Mosaic 是什么，以及为什么它是一条侧通道

Import 与 MosaicSerde 往返

导入 module 如何连接到 HLO

tpu-Dialect Pass Pipeline：RunMLIRPasses

TensorCore vs SparseCore：两条 Lowering 分支

TensorCore 分支（阶段 14–16）

SparseCore 分支 {#the-sparsecore-arm}

Mosaic 子页面地图

置信度总结

交叉引用