布局分配
地址适用于 libtpu-0.0.40-cp314 wheel 中的 libtpu.so。其他版本会有所不同。
摘要
布局分配是 XLA pass，它会在图被融合并 lowering 之前，固定模块中每个张量的物理 minor-to-major 维度顺序和 HBM tiling。在 TPU 上，这个 pass 是 xla::jellyfish::TpuLayoutAssignment（RunImpl @ 0x110ace00），它是开源 xla::LayoutAssignment 的子类。基类拥有双向约束传播工作列表；TPU 子类只覆盖后端扩展点：要种下哪些种子约束（AddBackendConstraints @ 0x110b19a0）、每个 op 如何为操作数/结果选择布局、哪些 op 可以改变布局，以及一个按张量的 fallback 选择器，用于传播永远到达不了的缓冲区（FindMemoryMinimizingLayout @ 0x1109dfe0）。它运行在 Phase 4（HloOptimizeThroughLayoutAssignment）内部，在 sharding/SPMD 之后、fusion 之前；参见 compile-phases.md。
熟悉上游 XLA LayoutAssignment 的读者会认出其骨架：种下强制约束，在 BufferLayoutConstraint/OperandLayoutConstraint 记录的工作列表上运行 forward+backward fixpoint，然后提交布局，并在生产者布局无法满足消费者所需布局时插入 kCopy。TPU 特有的是约束的内容：MXU 的 128-lane / 8-sublane tile 会驱动每个 dot、convolution、reduce-window、scatter 和 collective op 走向一种布局，使其两个最 minor 的维度能干净地铺进 tile；以及包在 OSS pass 外的一层重跑循环：RunImpl 会运行 OSS 布局分配直到 fixpoint，用 ModuleLayoutIsValid（0x110a6b80）验证，如果任何 TPU validator（Gather/Scatter/SelectAndScatter/ReduceWindow）拒绝，则重新播种并再次运行。
本页记录了重实现者必须复现的三件事：按张量的布局选择器 FindMemoryMinimizingLayout 及其真实成本模型（先看字节大小，再用 minor 维度填充效率比率打破平局，不是离散 penalty 的加权和）、AddBackendConstraints 按 opcode 种下的种子约束，以及这些选择如何通过 memory_space 整数和印在每个 xla::Layout 上的 tile 几何信息到达 memory-space assignment（MSA）与调度器。
对重实现而言，契约是：
- 驱动器形态。 TPU 重跑循环下的 OSS pass，并由 TPU validator 对每次迭代把关。
- 选择器。
FindMemoryMinimizingLayout：它的 element-type gate、rank>3 search gate，以及精确的双键成本：compact bytes，然后是 minor-dim fill ratio。 - 种子约束。
AddBackendConstraints中的 opcode→handler dispatch，以及每个 handler 强制执行的 MXU / packed-collective / window-stride==1 规则。 - 下游契约。 MSA 和调度从已分配的
xla::Layout中读取的内容（memory space、tile）。
| Pass 类 | xla::jellyfish::TpuLayoutAssignment : xla::LayoutAssignment : HloPassInterface |
| 驱动器 | TpuLayoutAssignment::RunImpl @ 0x110ace00 (0x37c4 ≈ 14 KB) |
| 种子约束 | TpuLayoutAssignment::AddBackendConstraints @ 0x110b19a0 (0x546d ≈ 21 KB; ModuleLayoutIsValid @ 0x585a ≈ 22 KB 略大) |
| 按张量选择器 | (anon)::FindMemoryMinimizingLayout @ 0x1109dfe0 |
| Compact-layout 选择器 | Target::ChooseCompactLayoutForShape @ 0x1d61bd00 |
| OSS 基类 RunImpl / AssignLayouts | 0x169bf440 / 0x169bb0e0 |
| 流水线位置 | Phase 4 尾部（HloOptimizeThroughLayoutAssignment），在 SPMD 之后、fusion 之前 |
| IR 层级 | XLA HLO（HloModule / HloInstruction），MLIR 之前 |
| 源文件 | platforms/xla/service/jellyfish/tpu_layout_assignment.cc (rodata) |
| MXU tile | (SublaneCount, LaneCount) = (8, 128) 此构建（v5+/ghostlite v6）；pufferfish v4 上为 (16, 128) |
| 置信度 | 除非某行或 callout 另有说明，均为 CONFIRMED（字节锚定） |
驱动器 — RunImpl 和重跑循环
目的
RunImpl 是该 pass 入口点的 TPU 覆盖实现。它不重新实现布局分配；它用设置逻辑、Mosaic kernel 解析步骤，以及外层 validate-and-re-run 循环包装开源 xla::LayoutAssignment::RunImpl（0x169bf440）。这层包装存在的原因是 OSS 算法分配的是对任意后端合法的布局；其中有些布局对 TPU 的 gather/scatter/reduce-window 限制并不合法，而发现这一点的唯一方式是运行 OSS 到 fixpoint、验证、再为违规者重新播种。
入口点
text
TpuLayoutAssignment::RunImpl (0x110ace00) ── driver, timed "XLA::JF Layout Assignment"
├─ GetTpuCompEnv / PartitionReplicaMapper setup ── multi-slice (megascale) topology
├─ PreprocessModule (0x1109ff20) ── insert kCopy boundary nodes
├─ ParseCustomCallKernelsInParallel ── pin Mosaic kernel I/O layouts
└─ loop:
├─ xla::LayoutAssignment::RunImpl (0x169bf440) ── OSS: seed → propagate → assign
│ └─ TpuLayoutAssignment::AddBackendConstraints ── (TPU vtable hook, fires inside OSS)
└─ ModuleLayoutIsValid (0x110a6b80) ── TPU validator; re-seed + repeat if false
```text
### 算法
```c
function RunImpl(module, exec_threads): // 0x110ace00
this->tpu_comp_env = GetTpuCompEnv(module) // member +1472 bytes (offset 184)
if this->target->GetMultiSliceTopology(): // megascale / ICI mesh
CHECK(module.config().static_device_assignment().has_value()) // src:3618
this->partition_replica_mapper =
PartitionReplicaMapper::Create(this->target, module.device_assignment())
ScopedLoggingTimer timer(name(), kXlaTimerTraceStats0) // guard 0x224cdec0
TraceMe("XLA::JF Layout Assignment")
// Insert kCopy boundary nodes so propagation has free choices at the seams.
PreprocessModule(module, exec_threads) // 0x1109ff20
// Mosaic (Pallas) kernels arrive with layouts pinned in their kernel proto.
// Parse them and stash per-operand / per-result layouts in the proposed maps.
ParseCustomCallKernelsInParallel(module) // populates members +0x420 / +0x440
for (instr, kernel) in tpu_custom_call instructions:
operand_layouts_map[instr] = CreateTwoMinorLayout(kernel.input_dims) // 0x110c1fa0
result_layouts_map[instr] = CreateTwoMinorLayout(kernel.output_dims)
passes_run = 0
while true:
// VLOG(2) "Before TpuLayoutAssignment: " ... src:3625
LogNeighborhoodFingerprints(module, "layout", 6, 6) // src:3626
// Reject ops that must have been removed by earlier expanders.
for c in module.MakeNonfusionComputationsSorted(exec_threads):
for inst in c.MakeInstructionPostOrder():
if inst.opcode in {batch-norm-grad/inf/training} or
(inst.opcode == fusion and inst.fusion_kind != kCustom):
return InvalidArgument(
"Instruction %s is not expected to be seen during layout assignment")
// The OSS engine: AddBackendConstraints (our hook) → PropagateConstraints
// → AssignLayouts → PostProcess.
xla::LayoutAssignment::RunImpl(module, exec_threads) // 0x169bf440 ; src:3737
++this->first_pass_count // member +183
if passes_run == 0: VLOG(2) "Ran first pass of layout assignment."
if not this->ModuleLayoutIsValid(module, exec_threads): // 0x110a6b80
VLOG(2) "Running " passes_run " additional passes of layout assignment ..."
// VLOG(3) the four offender lists (at-risk / scatter / scatter-updates /
// incompatible / expensive-collective), src:3679-3723
++passes_run
continue // re-seed: the recorded incompatibilities steer the next pass
break
// VLOG(2) "After TpuLayoutAssignment: " ... src post-loop
return /*module_changed=*/ true注意 — 重跑由验证失败驱动，而不是固定迭代次数。每次 OSS 运行本身都是一次完整的 forward/backward fixpoint；外层循环只会在 TPU 专属 validator 拒绝某个 OSS 引擎认为已完成的 op 时再次触发。实践中第二轮很少见：它会在某个 gather/scatter/reduce-window op 收到 TPU lowering 无法实现的布局时触发，并由记录下来的“incompatible layouts”列表引导下一轮重新播种。
函数映射
| 函数 | 地址 | 作用 |
|---|---|---|
TpuLayoutAssignment::RunImpl | 0x110ace00 | 驱动器 + 重跑循环 |
TpuLayoutAssignment::PreprocessModule | 0x1109ff20 | 分配前插入 kCopy 边界节点 |
TpuLayoutAssignment::ModuleLayoutIsValid | 0x110a6b80 | TPU 模块级 validator（循环 gate） |
xla::LayoutAssignment::RunImpl (base) | 0x169bf440 | OSS seed→propagate→assign 引擎 |
xla::LayoutAssignment::AssignLayouts (base) | 0x169bb0e0 | 将选定布局提交到模块 |
xla::LayoutAssignment::PropagateConstraints (base) | 0x169b8120 | Forward+backward 工作列表 |
(anon)::CreateTwoMinorLayout | 0x110c1fa0 | 从固定的 2-minor dims 构建 Layout |
GetReduceLayoutFromOperand (anon) | 0x110ac4e0 | 转发 reduce 布局 |
MemorySpaceColorMap::UpdateFromLayout | 0x110411a0 | 将 memory_space 读入按缓冲区的 color map（供 MSA 使用） |
特性 —
TpuLayoutAssignment成员在反编译中以*((qword*)this + N)访问。最重要的两个是：target_位于 qword+181（1448 字节处），每个 handler 都通过它访问芯片描述符；allow_relayoutflag 位于 qword+152（offset0x4C0），在任何 handler 被允许插入 relayout copy 之前都会测试*((_QWORD*)this + 152) != 0。
按张量选择器 — FindMemoryMinimizingLayout
目的
当约束传播永远到达不了某个缓冲区时，仍然必须为它选择布局；这通常发生在长 elementwise 链深处的浮点中间值上。对于 rank>3 张量，这个选择由 FindMemoryMinimizingLayout（0x1109dfe0）完成：它搜索 major 维度的排列，保留 compact byte size 最小的排列；如有平局，则按两个最 minor 维度填满 lane/sublane tile 的程度打破。它有两条到达路径：来自 GetUnconstrainedLayout（0x110b7520，OSS “没有约束到达此缓冲区” fallback），以及 AddBackendConstraints 在 copy-into-SparseCore 路径上的直接调用。
算法
c
function FindMemoryMinimizingLayout(target, shape /*in/out*/, out_bytes, cap): // 0x1109dfe0
// ---- Guard 1: element-type must be in the searchable set ----
// mask = 0x2FFF91FFE ; only proceed if primitive-type bit is set. movabs @0x1109e014
et = shape.element_type()
if et > 0x21 or not bittest64(0x2FFF91FFE, et):
return // type not layout-searchable; keep incoming
CHECK(shape.IsArray()); CHECK(LayoutUtil::HasLayout(shape)) // src:528-529
// ---- Baseline: compact byte size of the incoming layout ----
best_bytes = target->ShapeSizeCompact(shape) // 0x1d61a620
if out_bytes != null: *out_bytes = best_bytes
// ---- Guard 2: only search when rank > 3 ----
if shape.dimensions_size() <= 3: // @0x1109e0d8 (cmpq $3 ; ja)
return // low rank: keep OSS default layout
// ---- Minor-dimension tile fill factors (used only for the tie-break) ----
lane_chunks = shape.dimensions_minor(0) / target->LaneCount() // 0x1d60f400
sublane_chunks = shape.dimensions_minor(1) / target->SublaneCount() // 0x1d60f300
// ---- Search candidate permutations of the (rank>>1) major dims ----
for perm in candidate_major_permutations(shape):
apply perm to shape.layout().minor_to_major()
cand_bytes = target->ShapeSizeCompact(shape) // @0x1109e504
if cap.has_value() and cand_bytes > cap: continue
// KEY 1 (primary): strictly fewer bytes wins. jl @0x1109e77c (accept)
if cand_bytes < best_bytes:
best = perm; best_bytes = cand_bytes; continue
// KEY 2 (tie-break, equal bytes): lower fill-waste wins.
// cost = (per-dim utilization fractions, built via vcvtsi2sd/vdivsd/vmulsd)
// @0x1109e262 and @0x1109e73a, compared @0x1109e78c (vucomisd ; jbe @0x1109e790 reject)
if cand_bytes == best_bytes and cand_cost < best_cost:
best = perm; best_cost = cand_cost
commit best to shape.layout()
```text
### 成本模型 — 它是什么，以及不是什么
成本是**按优先级排列的两个 key**：
| Key | 数量 | 来源 | 规则 |
|---|---|---|---|
| 1 (primary) | Compact byte size `ShapeSizeCompact(shape)` | `0x1d61a620`（target-aware：tile + dtype packing） | 字节数严格更少 ⇒ 接受 |
| 2 (tie-break) | Minor-dimension fill-efficiency ratio（各维利用率分数的乘积） | 由 `dimensions_minor(0)/LaneCount` 和 `dimensions_minor(1)/SublaneCount` 构建 | 字节数相等时，fill-waste 更低 ⇒ 接受 |
函数中唯一的数值常量是 element-type 有效性 bitmask `0x2FFF91FFE`（bit *i* 置位 ⇒ primitive type *i* 可搜索）和 lane/sublane 除数（从芯片描述符读取，不是字面量）。反编译直接确认了二者：入口处的 `if ( v10 > 0x21 || (v11 = 0x2FFF91FFELL, !_bittest64(&v11, v10)) )`，以及从两个 minor-dim chunk count 构建比率的 `vcvtsi2sd xmm0,…,r12 / vcvtsi2sd xmm1,…,r15 / vdivsd` 序列。
> **注意 —** `FindMemoryMinimizingLayout` 内部没有 copy-count feedback：成本纯粹是最小 compact bytes，再用 minor-dim fill ratio 打破平局，没有离散 transpose/copy/bitcast penalty，也没有 `adaptive_layout_map_` 项。确实存在的 copy-count adaptivity 位于别处，即 `AdaptiveHloLayoutMap::RemoveOnCopyOverhead` @ `0x110accc0`，它作用于 pass *之间* profile-loaded overrides，而不是此选择器内部。
>
> **陷阱 —** rank gate 是 `> 3`，不是 `>= 3`。Rank-1/2/3 张量永远不会被搜索；它们保留 OSS `GetDefaultLayoutForRank` 布局。若重实现从 rank 3 开始搜索，将会生成与 libtpu 不同（且更慢收敛）的布局，尤其影响非常常见的 3-D activation shapes。
### 函数映射
| 函数 | 地址 | 作用 |
|---|---|---|
| `(anon)::FindMemoryMinimizingLayout` | `0x1109dfe0` | Rank>3 无约束选择器 |
| `TpuLayoutAssignment::GetUnconstrainedLayout` | `0x110b7520` | OSS fallback hook；调用选择器后再调用 `ChooseCompactLayoutForShape` |
| `Target::ShapeSizeCompact` | `0x1d61a620` | Compact byte size（成本 key 1） |
| `Target::LaneCount` / `SublaneCount` | `0x1d60f400` / `0x1d60f300` | Tile 除数（成本 key 2） |
| `Target::ChooseCompactLayoutForShape` | `0x1d61bd00` | 兜底 compact/SparseCore 选择器 |
---
## 种子约束 — AddBackendConstraints
### 目的
`AddBackendConstraints`（`0x110b19a0`）是 OSS 引擎在传播之前调用的 TPU vtable hook。它种下工作列表随后展开的*种子*约束。函数主体是一系列彼此独立的 `for inst in computation->MakeInstructionPostOrder()` 循环，每个循环都由 `HloInstruction+0xc` 上的内联 opcode-byte 测试把关（已验证：`movzbl 0xc(%r12),%eax`）。每个循环通过 OSS `SetInstructionLayout` / `SetArrayOperandLayout` / `SetOperandLayout`（`0x169b0740` / `0x169b00e0` / `0x169af600`）种下布局。
### 算法 — opcode dispatch
```c
function AddBackendConstraints(constraints): // 0x110b19a0
target = this->target_ // member qword +181
allow_relayout = (this->[qword 152] != 0) // offset 0x4C0
// LOOP 1 — HBM-transfer / RNG boundary ops -> compact tiled layout, tiles cleared.
for inst where opcode in {recv(89), rng-bit-generator(100), send(110)}:
L = target->ChooseCompactLayoutForShape(inst.operand(0).shape()) // 0x1d61bd00
ClearTiles(L); SetInstructionLayout(inst, L); SetArrayOperandLayout(inst.operand(0), L)
// LOOP 2 — adaptive / autofdo-proposed overrides (seeded before per-op loops).
for inst where this->adaptive_layout_map.HasLayout(inst): // 0x110b6e20
SetInstructionLayout(inst, adaptive_layout_map.GetLayout(inst), pin=false, set_default=false)
// LOOP 3 — Mosaic custom-call operand/result pinning (from RunImpl parse step).
for inst in operand_layouts_map (members +0x420 / +0x440):
if single-output: pin operand-0 layout = front()
elif tuple kCall(27)/kFusion(61): pin each tuple element's layout
// Collective ops -> packed-lane layout, replicas must agree.
for inst where CollectiveOpsNeedPackedLayout(inst): // 0x1109ef40
// switch on opcode: all-gather(6), all-gather-start(8), all-reduce(9),
// all-reduce-start(11), all-to-all(12), collective-broadcast(33),
// collective-permute(34), ragged-all-to-all(86), reduce-scatter(93)
SetInstructionLayout(inst, compact, pin=1, allow=1, default=1)
// LOOP 4 — the bulk: per-op layout choosers.
GetLayoutConfig(span, &conv_decisions /*this+0x500*/, &op_decisions /*this+0x560*/) // 0x110a5940
for inst in computation->MakeInstructionPostOrder():
switch inst.opcode:
case dot(52):
case convolution(43):
AssignConvolutionLayout(inst, this, target, &conv_decisions, allow_relayout) // 0x11096160
case reshape(97):
AssignReshapeLayout(target, inst, this, &op_decisions, allow_relayout) // 0x1109be80
case select-and-scatter(109):
AssignSelectAndScatterLayout(target, inst, this) // 0x1109b400
case ragged-dot(87):
AssignRaggedDotLayout(inst, this, target) // 0x1109a060
case reduce-window(94):
// max window dim over operand; if 2nd-minor extent >= SublaneCount() (runtime call)
// allocate new optimal layout; window bound/stride must be 1 in 2 minor dims
case concatenate(39):
CreateShapeWithOptimalLayoutForConcat(target, inst); pin concat layout // 0x110b1320
case gather(62):
GatherLayoutIsValid(inst, false) // 0x110a2d80
AddIndicesLayoutConstraintForScatterGather(inst, idx=dnums+0x90, this) // 0x110b7080
case scatter(107):
ScatterLayoutIsValid(inst) // 0x110a12a0
AddIndicesLayoutConstraintForScatterGather(inst, idx=dnums+0x90, this)
pin each update operand via GetScatterUpdatesLayout // 0x110a2600
case copy(44) feeding SparseCore:
if TransferSizeUtil::HasSparseCoreLayout(target.Topology(), op.shape()): // 0x110b7440
FindMemoryMinimizingLayout(target, op.shape(), ...); pin // 0x1109dfe0按类别划分的强制约束
上面的 handler 强制执行少量硬规则。重实现者必须复现的是这些规则；按 op 的 handler 只是机制。
| 类别 | Ops | 强制布局规则 | 执行者 |
|---|---|---|---|
| MXU feed | dot(52), convolution(43) | Contracting（input-feature）dim 作为内部 reduction 到达 MXU；output-feature/batch 通过 TwoMinorSize（0x110bd3a0）铺进 lane/sublane MXU tile | AssignConvolutionLayout (0x11096160) |
| HBM transfer | recv(89), send(110), rng-bit-generator(100) | Compact tiled layout，且 tiles cleared（传输边界处线性） | LOOP 1 inline |
| Collective | all-gather(6/8), all-reduce(9/11), all-to-all(12), collective-broadcast(33), collective-permute(34), ragged-all-to-all(86), reduce-scatter(93) | Packed-lane layout；所有 replica 一致 | CollectiveOpsNeedPackedLayout (0x1109ef40) |
| Gather/Scatter | gather(62), scatter(107) | Indices operand 固定在 index_vector_dim（packed）；scatter updates 通过 GetScatterUpdatesLayout 固定 | Add…ForScatterGather (0x110b7080) |
| Window | reduce-window(94), select-and-scatter(109) | Window bound 与 stride 在两个最 minor 维度中必须为 1；pass 会尝试让 window dim 外置 | AssignSelectAndScatterLayout (0x1109b400) + inline |
| Concat | concatenate(39) | 为所有 operands 选择单个 optimal layout | CreateShapeWithOptimalLayoutForConcat (0x110b1320) |
| Copy→SparseCore | copy(44) | Memory-minimising layout，然后固定 | FindMemoryMinimizingLayout (0x1109dfe0) |
陷阱 — window-op 规则（“bound and stride must be 1 in the two minor-most dims”）是运行时 fatal 的来源之一："Encountered select-and-scatter with a window bound or window stride in one of the two minor-most dimensions that are not 1, which is not implemented for TPU. XLA tries to choose a layout such that this is not the case, but that does not always succeed." 如果重实现不把 window dim 引向外侧，就会在合法 HLO 上触发它。
特性 — 没有单独的
kDot循环。Dot（opcode 52）和 convolution（opcode 43）共享同一个AssignConvolutionLayout调用点；AssignConvolutionLayout在入口处按0x34(dot)/0x2b(conv) 自分发。期望找到独立 matmul layout routine 的重实现者不会找到它：MXU tiling 数学是统一的。
函数映射
| 函数 | 地址 | 作用 |
|---|---|---|
TpuLayoutAssignment::AddBackendConstraints | 0x110b19a0 | 种子约束 dispatch（按 opcode 循环） |
TpuLayoutAssignment::AssignConvolutionLayout | 0x11096160 | dot+conv MXU 布局（自分发 0x34/0x2b） |
TpuLayoutAssignment::AssignReshapeLayout | 0x1109be80 | 通过 ImproveReshapeLayout 的 reshape 布局 |
TpuLayoutAssignment::AssignSelectAndScatterLayout | 0x1109b400 | Window-op 联合布局 |
TpuLayoutAssignment::AssignRaggedDotLayout | 0x1109a060 | RaggedDot（variable-batch matmul）布局 |
TpuLayoutAssignment::GatherLayoutIsValid | 0x110a2d80 | Gather validator |
TpuLayoutAssignment::ScatterLayoutIsValid | 0x110a12a0 | Scatter validator |
TpuLayoutAssignment::CollectiveOpsNeedPackedLayout | 0x1109ef40 | Collective packed-lane 测试（opcodes 6/8/9/11/12/33/34/86/93） |
TpuLayoutAssignment::GetLayoutConfig | 0x110a5940 | 将 autofdo proposed layouts 加载到 decision caches |
(anon)::AddIndicesLayoutConstraintForScatterGather | 0x110b7080 | 固定 indices operand layout |
(anon)::CreateShapeWithOptimalLayoutForConcat | 0x110b1320 | Concat optimal layout |
(anon)::TwoMinorSize | 0x110bd3a0 | Two-minor MXU tile sizing（conv 中有 8 个 call sites） |
传播 Hook — 按 op 的布局选择
目的
在每个 OSS 传播步骤内部，引擎会询问子类：“给定这个操作数布局，输出布局是什么？”（forward）以及“给定这个输出布局，操作数布局是什么？”（backward）。TPU 覆盖实现会为少数 opcode 回答，其余委托给 OSS 基类。它还覆盖两个策略谓词：哪些 op 可以改变布局，以及输出布局是否总是要推送到 operands。
算法 — forward / backward routing
c
function ChooseOutputLayoutFromOperandLayout(layout, inst, opnd_idx): // 0x110ba2c0 (forward)
switch inst.opcode:
case reduce(91): return GetReduceLayoutFromOperand(layout, inst) // 0x110ac4e0
case gather(62): return GetGatherOutputLayout(inst, layout, idx) // 0x110a4be0
case reshape(97): return ImproveReshapeLayout(layout, shape, …, target)// 0x110b8f00
default: return OSS::ChooseOutputLayoutFromOperandLayout(...) // 0x169b76a0
function ChooseOperandLayoutFromOutputLayout(layout, inst, opnd_idx): // 0x110b7de0 (backward)
switch inst.opcode:
case scatter(107): if opnd is the updates operand:
return GetScatterUpdatesLayout(target, inst, layout,
contiguous_update_window_dim) // 0x110a2600
case reduce(91): return <reduce-operand layout>
case reshape(97): return ImproveReshapeLayout(...) // 0x110b8f00
case fusion(61): return FindOperandPilots(...) // aggressive-loop-fusion path; flag-gated
default: return OSS::ChooseOperandLayoutFromOutputLayout(...) // 0x169b64c0
function OutputLayoutAlwaysPropagateToOperands(inst): // 0x11094ae0
if inst.opcode == 127 /*transpose*/: return true
return OSS::OutputLayoutAlwaysPropagateToOperands(inst) // 0x169b7660
function InstructionCanChangeLayoutDeepsea(inst): // 0x110b0640
switch inst.opcode:
case fusion(61): can-change iff fusion_kind == kCustom (else cmps 40/120/49)
case custom-call(49): look up CompilationProperties in the global custom-call
registry (0x10a87cc0); use its instruction_can_change_layout
case batch-norm-grad/inf/training 21/22/23: FATAL "!IsInvalidInstructionDuringLayoutAssignment" (src:3811)
default: return OSS::InstructionCanChangeLayout(inst) // 0x169c19c0
```text
> **特性 —** `OutputLayoutAlwaysPropagateToOperands` 对 **transpose（opcode 127）** 返回 true，这让 transpose 可以被实现为*纯布局变化*（免费），而不是数据 copy：所需的输出布局被强制推到 transpose 的 operand 上，因此 transpose 变成对 `minor_to_major` 的重新标记。反编译非常明确：`if ( *((_BYTE *)a2 + 12) == 127 ) return 1;`，而在 `HloOpcode` map 中 127 = `transpose`（该 enum 的最后一个值）。
>
> **注意 —** `InstructionCanChangeLayoutDeepsea` 的 custom-call 分支是该 pass 的可扩展性 hook。每个 TPU custom kernel（`tpu_custom_call`、`XlaMosaic`、`TopK`、`Sharding`、…）都会注册一个 `CompilationProperties` struct，其中的 `instruction_can_change_layout` boolean 会在这里被查询。重实现者若添加 custom op，必须注册它，否则 layout assignment 会按 OSS 默认方式处理它。
### 本页使用的权威 opcode 整数
上面的按 op dispatch 是针对编入此 libtpu 的上游 `HloOpcode` enum 的 byte test（从 `HloOpcodeString` @ `0x1e5ef000` 解码）。重实现者需要的整数如下：
| Opcode | Int | Opcode | Int | Opcode | Int |
|---|---|---|---|---|---|
| convolution | 43 (`0x2b`) | reduce | 91 (`0x5b`) | reduce-scatter | 93 (`0x5d`) |
| copy | 44 (`0x2c`) | reduce-window | 94 (`0x5e`) | gather | 62 (`0x3e`) |
| dot | 52 (`0x34`) | reshape | 97 (`0x61`) | scatter | 107 (`0x6b`) |
| custom-call | 49 (`0x31`) | fusion | 61 (`0x3d`) | select-and-scatter | 109 (`0x6d`) |
| collective-permute | 34 (`0x22`) | ragged-all-to-all | 86 (`0x56`) | ragged-dot | 87 (`0x57`) |
| concatenate | 39 (`0x27`) | transpose | 127 (`0x7f`) | sort | 120 (`0x78`) |
> **注意 —** `HloOpcode` enum 结束于 127（`transpose`）；没有 opcode 130。注意几个容易混淆的点：convolution 是 43（不是 49，49 是 custom-call），0x78 是 sort（不是 broadcast）。对于上表以外的任何 opcode integer，都应先重新推导再使用。
### 函数映射
| 函数 | 地址 | 作用 |
|---|---|---|
| `TpuLayoutAssignment::ChooseOutputLayoutFromOperandLayout` | `0x110ba2c0` | Forward 按 op routing |
| `TpuLayoutAssignment::ChooseOperandLayoutFromOutputLayout` | `0x110b7de0` | Backward 按 op routing |
| `TpuLayoutAssignment::OutputLayoutAlwaysPropagateToOperands` | `0x11094ae0` | transpose-always-propagate 规则 |
| `TpuLayoutAssignment::InstructionCanChangeLayoutDeepsea` | `0x110b0640` | 按 op 的 can-change bitmap |
| `(anon)::ImproveReshapeLayout` | `0x110b8f00` | Reshape source↔target dim mapping（双方向） |
| `(anon)::GetScatterUpdatesLayout` | `0x110a2600` | Scatter-updates backward layout |
| `(anon)::GetGatherOutputLayout` | `0x110a4be0` | Gather forward output layout |
| `xla::LayoutAssignment::InstructionCanChangeLayout` (base) | `0x169c19c0` | OSS fallback |
---
## 深度感知的布局成本与 Tile
### Tile 几何
每个已分配的 `xla::Layout` 都携带一个物理 HBM tile，其 dims 为 `(SublaneCount, LaneCount)`。这些值从 `target->[0x3b8]` 处的运行时芯片描述符读取，而不是硬编码：
- `Target::LaneCount()` (`0x1d60f400`) = `target->[+0x3b8]->[+0x198]` — 所有世代均为 `128`。
- `Target::SublaneCount()` (`0x1d60f300`) = `target->[+0x3b8]->[+0x1a0]` — 此 v5+/v6 构建上为 `8`，jellyfish v4 上为 `16`。
- `Target::ChunksPerTile()` (`0x1d60f2c0`) = `[+0x198] / [+0x1a0]`（lane_count / sublane_count，来自同一描述符）。
因此此构建的默认 tile 是 `(8, 128)`，v4 上是 `(16, 128)`。tile 和 `memory_space` 由 `HardwareLayout::PopulateDefaultLayout`（`0x1d6da120`）/ `HardwareLayout::PopulateShape`（`0x1d6da360`）印到 leaf subshapes 上。另一个按 shape 的 fixup `(anon)::UpdateLayout(target, Layout, Shape&)`（`0x110f66a0`）会把选定的 `Layout` 复制到 shape 上，然后委托给 `Target::UpdateLayout`（`0x1d618aa0`），后者运行 `TransferSizeUtil::UpdateLayout`（`0x1d6b05a0`）以在 dtype 改变后重新计算 tile/packing。若干 `jellyfish` pass 暴露自己的 `UpdateLayout(Shape*)` shim（`TpuSpmdPartitioner` `0x127a4100`、`TpuBFloat16Propagation` `0x110128e0`、`TpuBFloat16Normalization` `0x11012980`、`TpuReduceWindowRewriter` `0x109589c0` 等），以便在 pass 改变 shape 后刷新布局元数据；这些是彼此独立的 shim，都会调用 `Target`/`TransferSizeUtil` 布局路径，而不是共享同一个 helper。
### “深度”在哪里进入
`FindMemoryMinimizingLayout` 中的成本是按张量、tile-aware 的，而不是 graph-depth-aware 的：它会在每个候选 `minor_to_major` 下计算该张量的 compact bytes，并偏好两个 minor dims 最完整填充 `(SublaneCount, LaneCount)` tile 的布局。整个 *pass* 的“depth-aware”行为来自迭代重跑循环（更深的链会暴露更多 incompatibilities，从而为下一轮重新播种），以及 `ScatterLayoutIsValid` 用来确认第 2 minor dim 能干净打包进 sublane tile 的 `Compact2ndMinorRatio`（`0x1d61a4e0`）packing 检查。
> **特性 —** `ShapeSizeCompact` 已经把 tile 和 dtype packing 折进它返回的字节数里，所以“minimise bytes”和“fill the tile”并不是彼此独立的目标：浪费 tile lanes 的布局会 pad 出更多字节。fill-ratio tie-break 只在两个排列 pad 到完全相同的 byte total、但浪费在 lane dim 与 sublane dim 上的分布不同的时候才有意义。
### 函数映射
| 函数 | 地址 | 作用 |
|---|---|---|
| `Target::LaneCount` / `SublaneCount` / `ChunksPerTile` | `0x1d60f400` / `0x1d60f300` / `0x1d60f2c0` | 来自芯片描述符的 tile dims |
| `Target::Compact2ndMinorRatio` | `0x1d61a4e0` | 第 2 minor packing ratio（由 `ScatterLayoutIsValid` 调用） |
| `(anon)::UpdateLayout(target, Layout, Shape&)` | `0x110f66a0` | 按 shape fixup：设置 Layout，然后调用 `Target::UpdateLayout` |
| `Target::UpdateLayout` | `0x1d618aa0` | 通过 `TransferSizeUtil::UpdateLayout` 重新计算 tile/packing |
| `HardwareLayout::PopulateDefaultLayout` | `0x1d6da120` | 将 tile + memory_space 印到 leaf subshapes |
| `HardwareLayout::PopulateShape` | `0x1d6da360` | Tile-kind selection（Default/X64/X128） |
---
## 布局如何供给 MSA 和调度
### 目的
布局分配是整个后端的前提。它写入的两类元数据会被下游消费：嵌入每个 `xla::Layout` 的 `memory_space` 整数，以及 tile 几何信息。Memory-space assignment（MSA）读取前者来知道缓冲区位于哪个物理层级；调度器则按这些层级为 spill/refill 定价。
### 交接
`memory_space` 整数路由物理放置。`MemorySpaceColorMap::UpdateFromLayout`（`0x110411a0`）在布局分配后遍历模块中的每个 `Shape`，并将 `Layout::memory_space()` 读入按缓冲区的 color map；`BuildFromLayoutAndBackendConfig`（`0x11040d80`）还会合入 `frontend_attribute`/`backend_config` memory-space 字符串。这个 color map 就是 buffer allocator（MSA）消费的内容；参见 [msa-overview.md](msa-overview.md)。整数码由 `MemorySpaceToColor` 生成；命名空间包括 `kHbm`（off-chip）、`kVmem`（on-chip vector SRAM）、`kCmem`/`kSmem`（scalar/sequencer SRAM）、`kAlternate`（MSA 预留的 secondary HBM pool），以及 SparseCore/BarnaCore 空间。
Fusion（Phase 5）在布局分配*之后*运行，因为 fusion 合法性取决于物理 tile 布局；参见 [fusion-patterns.md](fusion-patterns.md)。调度器在 MSA 之后、位于 [scheduling pipeline](../sched/overview.md) 中运行，并按已分配的 memory spaces 为每个 op 以 cycles 定价；强制产生 HBM round-trip 的布局比将缓冲区保留在 VMEM 中的布局更昂贵，这正是此阶段最小化字节数（从而最小化 tile padding / HBM footprint）会影响最终 schedule 的原因。
> **注意 —** SparseCore 有自己的下游布局 pass：`xla::tpu::sparse_core::RunLayoutAssignment`（`0x12e02140`，src `sparse_core_layout_assignment.cc`），它在 post-layout pipeline 中运行于 SparseCore custom-call instructions 上（由 `sparse_core_compiler_util` 驱动，印上按 leaf 的 `HardwareLayout::TileKind`）。这里记录的 TensorCore 侧 `TpuLayoutAssignment` 会为 SparseCore 侧的 `memory_space` tags 播种（LOOP 1 / copy→SparseCore path via `HasSparseCoreLayout`），因此两个 pass 通过 layout 的 memory space 耦合。（此构建中命名的 SparseCore memory spaces 为 `kSparseCorePrivateStackHbm`、`kSparseCoreSequencerSmem`、`kSparseCoreSequencerSflag`，没有扁平的 `kSparseCoreMem` enumerator。）
---
## 相关组件
| 名称 | 关系 |
|---|---|
| `xla::LayoutAssignment` (OSS base, `0x169bf440`) | 拥有传播工作列表；TPU 类只覆盖 backend hooks |
| `TpuLayoutAssignment::PreprocessModule` (`0x1109ff20`) | Pre-pass：插入 `kCopy` 边界节点 |
| `MemorySpaceColorMap` (`0x110411a0`) | 读取为 MSA 分配的 `memory_space` |
| `xla::tpu::sparse_core::RunLayoutAssignment` (`0x12e02140`) | SparseCore 侧布局，位于下游，通过 memory space 耦合 |
| `TpuTilingRewriter::RewriteHostComputeWithLinearLayout` (`0x1112efc0`) | Post-layout：将 tiled TensorCore layouts 桥接到 host flat memory |
## 交叉引用
- [compile-phases.md](compile-phases.md) — 有序 phase 主干；布局分配是 Phase 4 尾部，在 SPMD 之后、fusion 之前
- [overview.md](overview.md) — Part V 定位；布局分配位于 XLA-then-MLIR 下降过程中的位置
- [hlo-pre-passes.md](hlo-pre-passes.md) — 必须先运行的 pre-passes（包括 BatchNormExpander）；layout assignment 会对未展开的 batch-norm FATAL
- [fusion-patterns.md](fusion-patterns.md) — 主 fusion（Phase 5），它在布局之后运行，因为 fusion 合法性取决于物理 tile 布局
- [msa-overview.md](msa-overview.md) — memory-space assignment，消费这里写入的 `memory_space` color map
- [mosaic-layout-inference.md](mosaic-layout-inference.md) — Mosaic/Pallas kernel layouts，在 `RunImpl` 中解析并固定为种子约束
- [../sched/overview.md](../sched/overview.md) — 调度器会根据此 pass 分配的 memory spaces 为 op 以 cycles 定价