循环平铺与展开

本页中的所有地址都适用于来自 libtpu-0.0.40-cp314 wheel 的 libtpu.so（build-id 89edbbe81c5b328a958fe628a9f2207d）。其他版本会有所不同。.text 地址是虚拟地址；对于此二进制文件，.text VMA == 文件偏移 0xe63c000。

摘要

TPU 编译器中没有单一的“循环优化器”。来自 LLVM 背景的读者所期待的那种机制——一个 pass 对循环嵌套做平铺，并按成本驱动的因子展开内层循环——在这里被拆分到三个互不相关的机制中，它们作用于三种不同的 IR，并由三个独立的标志控制。本页负责这套机制中属于循环形状决策的部分：TileKind 规则，它用 compact 或 SparseCore tile 标签标记每个 shape 的 layout（TpuTilingAssignment::GetCopyTileKind, 0x13dd0ca0）；LoopConfig / LoopUnrollConfig proto，它们携带 SparseCore kernel 收到的逐循环维度展开指令（逐字段解码）；以及 HLO while-loop pipeline unroller（WhileLoopPipelineUnroller, 0x12ee2200），它把可软件流水的 while 转换为 depth 个 stage 的链式 call body。

这里称为“tiling”的东西不是 GPU/Triton 的 tile size。在 TPU TensorCore 路径上，它是写入每个 Shape::Layout 的 16-bit TileKind 标签——对于 copy，是一对 (input_kind, output_kind)，其中每个 kind 取自 {Default=0, SparseCore=3}——后端其余部分读取它来判断 buffer 保持真实的 2-D tile，还是采用 SparseCore layout。真正按成本排序的 MXU operand-window tiling 位于 dot/conv → MXU lowering page；SparseCore tile-index 地址展开（ExpandTiledMemRefs）位于 Tile-Index Expansion。LLVM 侧为这些 pass 留下的循环寻找 initiation interval 的 modulo scheduler 位于 Bundle Modulo Scheduling。本页不重新推导其中任何内容；它只负责 TileKind 决策、LoopConfig schema，以及展开/流水转换。

对于重新实现，契约如下：

• TileKind 是 layout 标签，不是 tile size。 TpuTilingAssignment 在 layout assignment 后遍历 module，并用 TileKind 标记每个 shape 的 Layout。对于 kCopy，它从四个 TransferSizeUtil layout predicate 计算 packed (input,output) 对；其他所有 opcode 都通过 HardwareLayout::GetDefaultLayout 继承 layout-assignment tile。这条路径上没有丰富的逐 opcode tile-shape 求解器。
• SparseCore loop 携带的 unroll factor 位于 LoopConfig proto 中。 字段 3（unrolled_loops，repeated LoopUnrollConfig）用 loop_dim（字段 1）作为每个 entry 的 key，并携带显式 unroll_factor 或 “auto” sentinel（一个 oneof）。normalizer 通过将 loop bound 除以 vectorizing_shape（字段 4）的元素数来解析 “auto”，并带有 exact-divisibility CHECK。
• 逐架构 SparseCore copy unroll factor 是一个双 template switch。 CustomLoopUnrollPolicy<5>（template constant tpu::TpuVersion 5 = 6acc60406，营销名 “Ironwood”）vs <3>（template constant tpu::TpuVersion 3 = viperfish，作为所有非 6acc60406 版本的 fallback），按 Target::tpu_version 分派。<3> template 发出 16（elementwise）/ 8（structured）；<5> 发出 transpose {16,8}、general 32/16，或 MD-vectorizing 16/pack。
• While-loop 展开与软件流水是独立且可叠加的转换。 WhileLoopUnroller（full / double-buffer / auto）和 WhileLoopPipelineUnroller（loop-carry-depth pipelining）由两个独立 env knob 控制，并且可以同时运行在同一个 module 上。pipeliner 将 body 克隆成 depth 个链式 call stage，并把 trip count 减少 depth-1。

TileKind 标记（HLO）	xla::jellyfish::TpuTilingAssignment : HloPassInterface; name() returns "tiling-assignment"
TileKind 决策	TpuTilingAssignment::GetCopyTileKind @ 0x13dd0ca0（已反编译，按字节锚定）
TileKind 驱动器	VerifyOrAssignTiling @ 0x10922a20; RunImpl @ 0x13dd10a0; Verify @ 0x13dd2900
Post-fusion 特殊 tiling	TpuPostFusionTilingAssignment @ RunImpl 0x13dd85a0; name() returns "post-fusion-tiling-assignment"
LoopConfig proto	xla::jellyfish::LoopConfig（serializer 0x1d6eade0）；LoopUnrollConfig（serializer 0x1d6f2680）
展开算术	LoopConfigWrapper::GetNormalizedUnrollFactor @ 0x13d6c1c0（已反编译）
逐架构 SC copy factor	GetCustomLoopUnrollPolicy @ 0x13916ec0; CustomLoopUnrollPolicy<(tpu::TpuVersion)5>::GetConfig<(HloOpcode)44> @ 0x13916fe0; <(tpu::TpuVersion)3> @ 0x139173a0
While-loop unroller	xla::WhileLoopUnroller; IsLoopUnrollable @ 0x12ee8620; name() returns "while_loop_unroller"
Pipeline unroller	xla::WhileLoopPipelineUnroller::RunImpl @ 0x12ee2200; ComputeWhileLoopPipelineDepth @ 0x12ee0fc0; name() returns "while_loop_pipeline_unroller"
SC window selector	WindowUnrollFactorSelector::Select @ 0x1385c360; name() returns "window-unroll-factor-selector"
Pipeline 宿主	PostOptimizationPipeline @ 0x1093fd40（HLO loop passes）；RunBackendWithBufferAssignment @ 0x13070bc0（SC selector）
置信度	CONFIRMED（按字节锚定），除非某行或 callout 另有说明

---

三种“loop”机制，需要分开看待

“tiling”和“unrolling”这两个词在此二进制文件中各自指代三种不同的东西。重新实现者必须把它们分开，否则会构建错误的 solver。

关注点	机制（本页）	作用的 IR	决策风格
TileKind layout 标签	TpuTilingAssignment 按 shape 标记 (Default|SparseCore)	HLO Shape::Layout	正确性/layout predicate，不是成本模型
SparseCore loop unroll factor	LoopConfig.unrolled_loops（由 WindowUnrollFactorSelector + CustomLoopUnrollPolicy 写入）；由 MLIR LoopUnrollPass 应用	jellyfish HLO → SC MLIR scf.for	逐架构表 + scratchpad 受限的贪心 fit
HLO while-loop 展开 / 流水	WhileLoopUnroller, WhileLoopPipelineUnroller	HLO kWhile	trip-count/body-size 阈值；loop-carry depth

本页有意不负责的三件事，以及它们的位置：

• 按成本排序的 MXU operand-window tiling（IterateThroughWindowConfigs，按 MXU cycles 定价，受 VMEM 限制）——这是唯一真正成本驱动的 tiling 搜索，并且属于 convolution 路径。参见 Dot / Conv → MXU Lowering。
• SparseCore tile-index 地址展开（ExpandTiledMemRefs / expandTiledIndices）——逐 tile memref index algebra。参见 Tile-Index Expansion。
• bundle modulo scheduler——LLVM 侧的 initiation-interval（II）搜索，它对 SC loop unroller 留下的 LLO bundle stream 做软件流水。这不同于本页的 HLO while-loop pipelining。参见 Bundle Modulo Scheduling。

NOTE — TensorCore 分配的 “tiling” 是标签，不是大小。 HLO TpuTilingAssignment pass 不求解 tile shape。实际 tile geometry（Default 128×16 / X64 / X128）由 layout assignment 产生并存储在 Layout 中；此 pass 只写入 TileKind enum，说明 shape 使用哪一类 tile，因此 lowering legalizer 可以读回它。若重新实现者在这里构建 tile-size 搜索，那就是在构建错误的 pass。

---

TileKind：layout 标记

驱动器与模式（VerifyOrAssignTiling, 0x10922a20）

xla::jellyfish::DeepseaCompilerBase::VerifyOrAssignTiling(const Target&, HloModule*) 是决定是否以及如何提交 TileKind 标签的单一入口。它读取 TpuCompilationEnvironment + 0xDFC 处的 tri-state int（十进制 3580；标志 xla_tpu_verify_or_assign_tiling_before_lowering，另有两个带后缀变体 ...lowering3 / ...lowering8 作为独立 flag string 存在）：

absl::Status DeepseaCompilerBase::VerifyOrAssignTiling(
    const Target& target, HloModule* module) {
  int mode = *(int*)(GetTpuCompEnv(target) + 3580);            // env + 0xDFC

  if (mode == 1) {                                  // VERIFY
    TpuTilingAssignment pass(target, /*ctor_bool=*/false);     // ctor arg = 0
    return pass.Verify(module, /*exec_threads=*/{});           // deepsea_compiler_base.cc:3053
  }
  if (mode == 2) {                                  // ASSIGN
    TpuTilingAssignment pass(target, /*ctor_bool=*/true);      // ctor arg = 1
    return pass.Run(module, /*exec_threads=*/{}).status();     // deepsea_compiler_base.cc:3056
  }
  return absl::OkStatus();                          // mode == 0: leave tiling as-is
}
```text

`TpuTilingAssignment(const Target&, bool)` ctor（`0x13dd1080`）在内部选择 verify-vs-assign；反编译显示 `mode == 1` 用 `bool = 0` 构造并调用 `Verify()`，而 `mode == 2` 用 `bool = 1` 构造并调用 `Run()`。因此 ctor 的 `bool` *不是* “verify_only” 标志——`1` 是 assign/Run 路径。手动覆盖是一个 tri-state int：`0` skip，`1` verify，`2` assign。

### 逐 copy 的 TileKind 算法（`GetCopyTileKind`, `0x13dd0ca0`）

这是逐字节恢复出来的唯一 TileKind 规则。对于 `kCopy`，它返回一个 `StatusOr`，其中 packed 16-bit `TileKind` = `(input_kind | (output_kind << 8))`，每个 kind 都在 `{Default=0, SparseCore=3}` 中。*input* 侧读取 `operand(0)` 的 shape；*output* 侧读取 `copy.shape()`；两边运行完全相同的四 predicate 链。

```cpp
StatusOr<uint16_t> GetCopyTileKind(const HloInstruction& copy, const Target& target) {
  TransferSizeUtil* tu = target.transfer_size_util();   // Target + 0x3B8 (qword 119)
  uint16_t in_kind, out_kind;

  // --- INPUT side: operand(0) shape ---
  const Shape& in = copy.operand(0)->shape();
  int ms = in.layout().memory_space();                  // Shape + 0x138 (byte 312)
  const Shape& real_in = (ms == kVmem /*3*/ || ms == kSmem /*5*/)
                       ? in.tuple_element_or_self() : in;
  if (real_in.layout().has_minor_to_major()             // Shape + 0x130 (byte 304) == 1
      && in.layout().minor_to_major().size() >= 2       // layout()[9] >= 2
      && !tu->HasLinearLayout(in)
      && !tu->HasSparseCoreLayout(topo, in)
      && !tu->HasPadless2ndMinorLayout(topo, in)
      && !tu->HasLarge2ndMinorLayout(topo, in)) {
    in_kind = 0;                                        // Default (compact 2-D tile)
  } else if (tu->HasSparseCoreLayout(topo, in)) {
    in_kind = 3;                                        // SparseCore tile
  } else {
    return InvalidArgument(
        "Input shape does not have compact or sparse core layout.");  // :66
  }

  // --- OUTPUT side: copy.shape() (same chain) ---
  const Shape& out = copy.shape();
  if (/*…same four-predicate chain on out…*/) {
    out_kind = 0;
  } else if (tu->HasSparseCoreLayout(topo, out)) {
    out_kind = 0x300;                                   // = 3 << 8
  } else {
    return InvalidArgument(
        "Output shape does not have compact or sparse core layout.");  // :77
  }

  return in_kind | out_kind;   // stored at *((uint16_t*)result + 4); status word = OK(1)
}

结构逐字来自反编译（0x13dd0ca0）：input kind v18 ∈ {0,3}，output kind v23 ∈ {0,0x300}，结果 v18 | v23 写入 this+8。VMEM(3)/SMEM(5) operand 会在 has-bit 检查前解引用 tuple-element shape（v9 = *(v7+8)）。四个 TransferSizeUtil layout predicate 是 gate——它们决定某一侧保持真实 2-D tile（kind 0），还是 SparseCore（kind 3）：

Predicate	地址	含义
HasLinearLayout(Shape)	0x1d6af220	untiled/linear layout（host-transfer 边界）
HasSparseCoreLayout(TpuTopology, Shape)	0x110b7440	SparseCore tiling
HasPadless2ndMinorLayout(TpuTopology, Shape)	0x1d6af3e0	no-pad 2nd-minor
HasLarge2ndMinorLayout(TpuTopology, Shape)	0x1d6af2e0	“special” / large 2nd-minor

GOTCHA — copy TileKind 是一个对，并且 mixed pair 合法。 如果一个 kCopy 的 operand 是 compact，而 result 是 SparseCore-laid-out，则返回 0 | 0x300 = 0x0300；全 SparseCore 情况是 0x0303；全 compact 情况是 0x0000。若重新实现者把 TileKind 折叠为单个逐 shape 值，将无法 round-trip 跨 compact↔SparseCore 边界 re-tile 的 copy。

其他每个 opcode 都继承 layout-assignment tile

RunImpl（0x13dd10a0）不会为非 kCopy op 计算 (input,output) 对。它用 ShapeUtil::ForEachMutableSubshape（0x13dd26a0 / 0x13dd27e0 处的 $_4/$_5 visitor）遍历 result shape，并把 HardwareLayout::GetDefaultLayout(result_shape) 标记到每个 leaf 的 mutable_layout()。只有以下 opcode 有特殊处理：

• kCopy — 上述 (input,output) 对。
• kOutfeed（opcode 80）— 使用 mutable_outfeed_shape()。
• tuple-result ops（例如 kReduce 类）— 通过 subshape visitor 递归 tuple，标记每个 leaf 的 default tile。
• async-SparseCore（async_execution_thread == "sparsecore"）— 从 async-chain start 的 backend config 读取 SparseCoreConfig，并标记 HardwareLayout::FromProto(config)——T8 / SC-tiling opt-in 路径（IsT8CustomKernelInstruction, 0x13dd0e80）。

一个 dtype gate 保护 general stamp：element_type 必须在 supported-tiled mask 0x2FFF91FFE（或一个小的额外集合）中，否则 instruction 必须是 fused（CHECK instruction->IsFused(), tpu_tiling_assignment.cc:233）并被跳过。

NOTE — 没有丰富的逐 HLO opcode TileKind 分化。 唯一的“规则”就是上面四条；其他所有内容都通过 GetDefaultLayout 保留 layout-assignment 选择的 tile。这就是为什么 TileKind 被记录为标记，而不是 solver。

post-fusion pass TpuPostFusionTilingAssignment（RunImpl 0x13dd85a0）在很晚之后运行，位于 copy insertion 之后，并将 “special tiling”（HasLarge2ndMinorLayout family，由 xla_tpu_enable_large_2nd_minor_layout[_for_x{4,8,16}] 控制）从 CanProduceSpecialTiling（0x13dd7760）的 producer forward-propagate 到 AcceptsSpecialTiling（0x13dd6580）的 consumer，受 entry in/out set 和 alias constraint 约束。逐 opcode predicate set 已命名，但其 producer→consumer 规则没有逐行反编译（参见 Confidence table）。

---

LoopConfig / LoopUnrollConfig proto

SparseCore kernel 在每个 fusion 附加的 LoopConfig proto 中携带其 loop-tiling 和 unroll 指令。WindowUnrollFactorSelector 写入它；LoopConfigWrapper 读回它；MLIR LoopUnrollPass 应用它。两个 schema 都从其 _InternalSerialize wire emitter 逐字段解码（field number = tag byte >> 3；offset 是生成的 C++ message 中的 struct byte offset；has-bits 位于 _has_bits_ 中）。

xla::jellyfish::LoopConfig（serializer 0x1d6eade0）

Fld	Wire tag	类型	Struct off	名称（推断）	语义
1	0x0A	repeated int64（packed）	RepeatedField +0x18, cnt int +0x1C, data ptr +0x20; has-bit &1	loop_bounds	逐 dim trip / index space
2	0x10	int64	+0x58; has-bit &8	（scalar）	辅助 scalar（trip/total）
3	0x1A	repeated message	RepeatedPtrField +0x30, cnt int +0x38; has-bit &2	unrolled_loops	repeated LoopUnrollConfig
4	0x22	repeated int64（packed）	RepeatedField +0x40, cnt int +0x50, data ptr +0x48; has-bit &4	vectorizing_shape	逐 dim native vector shape

xla::jellyfish::LoopUnrollConfig（serializer 0x1d6f2680）

Fld	Wire tag	类型	Struct off	名称（推断）	语义
1	0x08	int64	+0x18	loop_dim	此 entry 作为 key 的 dim（join key）；has-bit _has_[0]&1
2	0x10	int64（存储 1 byte）	+0x30	auto_kind	“auto/full” sentinel — oneof case 2（共享 union slot）
3	0x18	int64	+0x30	unroll_factor	显式 factor — oneof case 3（共享 union slot）
4	0x20	bool	+0x28	pipeline_remainder	remainder-loop pipelining 标志；has-bit _has_[0]&4
5	0x28	int64	+0x20	（aux）	辅助值；has-bit _has_[0]&2

字段 2 和 3 组成一个 oneof：它们的 payload 共享 +0x30 处的同一个 union slot，discriminator oneof_case 是 +0x38 处的 dword（2=auto，3=explicit）。这被确认了两次：0x1d6f2680 处的 serializer 读取 case *((uint32*)this+14)（byte +0x38），并写入 field-2 byte +0x30 或 field-3 qword +0x30；而 GetLoopUnrollFactor（0x13d6c100）从 +0x38 读取 copied-out oneof_case，并从 +0x30 读取 payload。loop_dim（字段 1）是 join key：GetLoopUnrollConfig(dim)（0x13d6c080）线性扫描 unrolled_loops，寻找 field-1 == dim 的 entry。

NOTE — proto 字段名是从 wire-format 逆向工程推断出来的，不是 symbol string。 tag、type 和 offset 都由 serializer 按字节锚定；人类可读名称（loop_bounds, unrolled_loops, vectorizing_shape, loop_dim, unroll_factor）是结合周边代码和 CHECK string 后最一致的解读。名称为 MEDIUM confidence；layout 为 CONFIRMED。

---

Unroll-factor 算术

GetNormalizedUnrollFactor（0x13d6c1c0）

这是把某个 LoopUnrollConfig entry 转换为给定 loop bound 的具体 factor 的算法。反编译（0x13d6c1c0）：

StatusOr<int64_t> LoopConfigWrapper::GetNormalizedUnrollFactor(
    absl::Span<const long> bounds, long dim) const {
  const LoopUnrollConfig* cfg = GetLoopUnrollConfig(dim);   // scan unrolled_loops by loop_dim
  if (!cfg) return 1;                                       // no entry → factor 1

  int64_t raw = (cfg->oneof_case() == 2) ? (-(uint8_t)cfg->auto_kind() | 1)  // auto sentinel
              : (cfg->oneof_case() == 3) ?  cfg->unroll_factor()             // explicit
              :                              1;
  if (raw != -1) return raw;                  // explicit / non-auto factor is final

  // --- AUTO path (raw == -1): divide loop bound by the VECTORIZING SHAPE ---
  int vs_size = loop_config_.vectorizing_shape_size();      // field 4 count, env-of-this + 0x44
  if (vs_size == 0)
    return InvalidArgument("Vectorizing shape missing");                       // 25-char
  if (vs_size >= 2)
    return InvalidArgument(
        "Vectorizing shape has too many dimensions: %d", vs_size);            // 45-char fmt

  int64_t loop_bound       = bounds[dim];
  int     vectorizing_shape = loop_config_.vectorizing_shape(0)[8];  // the dim count
  CHECK(loop_bound % vectorizing_shape == 0)                  // loop_config_wrapper.cc:358
      << loop_bound << " % " << vectorizing_shape;
  return loop_bound / vectorizing_shape;
}
```text

反编译精确确认了这一点：对 `unrolled_loops`（位于 `a2+48`，count `a2+56`）按 `loop_dim`（`*v7+3`）匹配的线性扫描，oneof 读取（`v36 == 2` → `-(uint8)v35 | 1`；`== 3` → `v35`；否则 `1`），`!= -1` 短路，以及 auto path 读取 `vectorizing_shape_size`（`*(a2+68)`）、两个错误字符串、`bounds[dim]`、`vectorizing_shape(0)[8]`，和第 358 行带 `LogMessageFatal` 的 `loop_bound % vectorizing_shape == 0` CHECK。

> **GOTCHA — “auto” divisor 是 `vectorizing_shape`，不是 unroll-config count。** 对于标记为 `auto` 的 loop，normalized factor 是 `loop_bound / vectorizing_shape[0].dim_count`，并要求 exact divisibility。若重新实现者除以 `unrolled_loops_size`（一个看似合理的误读），将产生错误 factor，并漏掉 `"loop_bound % vectorizing_shape == 0"` invariant。支撑 helper：`GetLoopUnrollFactor`（`0x13d6c100`）返回 raw factor（oneof 3 → `unroll_factor`；oneof 2 → `-(uint8)auto_kind`）；`GetLoopPipelineRemainder`（`0x13d6c4e0`）在设置时返回 `(remainder_byte | 0x100)`，默认 `0x101`。

### 逐架构 SparseCore copy policy

`GetCustomLoopUnrollPolicy(SmallVector<long,6> bounds, HloInstruction, Target)`（`0x13916ec0`）按 `Target::tpu_version`（`Target + 0x398`）分派：

```cpp
LoopUnrollPolicy GetCustomLoopUnrollPolicy(const SmallVector<long,6>& bounds,
                                           const HloInstruction& hlo,
                                           const Target& target) {
  int version = target.tpu_version();                       // Target + 0x398 (920)
  const HloInstruction* copy = lowering_util::GetCopyInstruction(hlo);
  if (!copy) return {};                                     // empty policy
  LoopUnrollPolicy p =
      (version == 5)                                        // tpu::TpuVersion 5 = 6acc60406 ("Ironwood")
        ? CustomLoopUnrollPolicy<5>::GetConfig<kCopy>(bounds, *copy, target)
        : CustomLoopUnrollPolicy<3>::GetConfig<kCopy>(bounds, *copy, target);  // every other version
  // CHECK each returned unroll_dimension ∈ [0, bounds.size())  (perf_utils.cc:151/152)
  return p;
}

C++ tpu::TpuVersion enum 是 proto enum 减一——kJellyfish=0, kDragonfish=1, kPufferfish=2, kViperfish=3, kGhostlite=4, k6acc60406=5——因此 dispatch literal 5 是 6acc60406（TPU_VERSION_* proto descriptor 将这些编号为 1..6，确认 −1 offset）。只有两个 arch template：<(tpu::TpuVersion)5>（6acc60406，营销名 “Ironwood”），仅在 version == 5 时采用；以及 <(tpu::TpuVersion)3>（template constant viperfish），作为所有其他版本（jellyfish 到 ghostlite）的 fallback。

CustomLoopUnrollPolicy<(tpu::TpuVersion)3>::GetConfig<kCopy>（0x139173a0，fallback template） — 已反编译，按字节锚定：

LoopUnrollPolicy CustomLoopUnrollPolicy</*tpu::TpuVersion*/3>::GetConfig<kCopy>(
    const SmallVector<long,6>& bounds, const HloInstruction& hlo, const Target&) {
  CHECK(hlo.opcode() == HloOpcode::kCopy);                  // perf_utils.cc:43
  int inner = bounds.back();
  bool elementwise = lowering_util::IsElementwiseCopy(hlo);
  // ONE entry: dim = inner - 1, factor = 8 * elementwise + 8
  return { { /*dim=*/inner - 1, /*factor=*/ (elementwise ? 16 : 8) } };
}
```text

反编译显示 `*(v5+8) = 8 * IsElementwiseCopy + 8` 且 `*(v5) = v3 - 1`——也就是在最内层 dim 上，elementwise copy 的 factor 为 **16**，structured copy 为 **8**。

**`CustomLoopUnrollPolicy<(tpu::TpuVersion)5>::GetConfig<kCopy>`（`0x13916fe0`，6acc60406 / “Ironwood”）** — 逐字节反编译：

```cpp
LoopUnrollPolicy CustomLoopUnrollPolicy</*tpu::TpuVersion*/5>::GetConfig<kCopy>(
    const SmallVector<long,6>& bounds, const HloInstruction& hlo, const Target& target) {
  CHECK(hlo.opcode() == HloOpcode::kCopy);                  // perf_utils.cc:76
  int rank = bounds.size(); int inner = bounds.back();
  PrimitiveType dtype = hlo.shape().element_type();
  // dtype acceptance mask 0x2FFF91FFE ∪ {0x20,0x21,15,18}; else FATAL (primitive_util.h:757)

  // TRANSPOSE-COPY special case: sub-word packed dtype (mask 0x910) &&
  //   IsMinorTransposeCopy(hlo) && rank >= 2 (CHECK "rank > 1", perf_utils.cc:81) && inner < 32
  if (sub_word_packed(dtype) && lowering_util::IsMinorTransposeCopy(hlo)
      && rank >= 2 && inner < 32) {
    return { {rank-1, 16}, {rank-2, 8} };                   // TWO entries
  }

  // GENERAL case
  bool pred1 = TransferSizeUtil::ShouldPackPREDAsSingleBit(topo, hlo.shape());
  int  pack  = TransferSizeUtil::ElementPackingFactor(dtype, pred1);
  CHECK(target.SupportsSparseCore());                       // target.h:1709
  int scs_tc = target.topology()->sc_tile_count();          // topology + 148
  int subl   = target.SublaneCount();
  LoopConfigWrapper w = LoopConfigWrapper::Create(hlo, rank, scs_tc, pack, subl);
  int factor = lowering_util::IsMDVectorizingShape(pack, target, /*…*/)
             ? 16 / pack
             : ((inner % (32 * scs_tc)) == 0 ? 32 : 16);
  return { {rank-1, factor} };                              // ONE entry
}

Template	Case	Unroll factor
<3>（fallback：所有 version != 6acc60406）	elementwise copy	16
<3>	structured copy	8
<5>（6acc60406 / “Ironwood”）	narrow transpose copy（inner < 32）	{16, 8}（两个 dim）
<5>	MD-vectorizing shape	16 / pack
<5>	general, inner % (32·scs_tc) == 0	32
<5>	general, otherwise	16

NOTE — template constant 不是营销芯片名。 CustomLoopUnrollPolicy<5> 是 C++ template parameter tpu::TpuVersion 5 = 6acc60406（proto descriptor 将其编号为 TPU_VERSION_6acc60406 = 6，所以 C++ enum 是 proto−1）；<3> 是 viperfish constant，用作 dispatch 未路由到 <5> 的所有版本的 catch-all。两个 GetConfig<kCopy> body 都逐字节反编译：<3> 位于 0x139173a0（8 * IsElementwiseCopy + 8 → 16/8），<5> 位于 0x13916fe0（transpose {16,8}；general 16 * (inner % (32·scs_tc) == 0) + 16 → 32/16；MD-vectorizing 16 / pack）。CONFIRMED。

---

SparseCore window unroll：受 scratchpad 限制的贪心

WindowUnrollFactorSelector 是选择 gather/scatter window unroll factor，并将 CustomLoopUnrollPolicy 结果写入 LoopConfig.unrolled_loops proto 的 HLO pass。它由 RunBackendWithBufferAssignment（0x13070bc0）添加到 SparseCore backend 自己的 HLO sub-pipeline 中，形式为 AddPass<WindowUnrollFactorSelector, Target const*, long>，其中 long 是 FLAGS_xla_sc_tiles；它是“lowering 前立即运行的两个 late annotation pass”之一（CHECK pipeline.PassesSize() == 2, sparse_core_compiler.cc:599）。

Select(instr, bool)（0x1385c360）：

1. 识别 gather-offload / scatter-offload custom fusion；对于 offloaded op，记录 "But this is an offloaded op. So, we will not find an unroll factor." 并退出（无 factor）。
2. 提取内部 kGather / kScatter 并分类 access pattern（IsSublane/IsElement/IsLane/IsChunk Gather/Scatter）。
3. 按模式读取 SC scratchpad budget：tile → MaxTileScratchpadSizeInBytes；SCS → MaxScsScratchpadSizeInBytes；loop-fusion → FusionEmitter::GetReservedScratchpadBytes（或 GetReservedSmemBytes）。
4. 选择能 fit 的最大 candidate factor——这是贪心 resource fit，不是 roofline cost。
5. 对每个返回的 CustomLoopUnrollPolicy entry，构建一个 LoopUnrollConfig（field-1 = loop_dim，oneof-3 = unroll_factor），并将其 Add 到父 LoopConfig.unrolled_loops。

逐 candidate fit test，ChunkGatherWindowSizeFitsInScratchpad(target, instr, factor)（0x1385c240）：

window_elems = Product(GetSliceSizesTiledPadding(instr))
bytes        = window_elems * ByteSizeOfPrimitiveType(dtype)
sized        = 8 * ((bytes >> 2) + 1) * factor   // round up to 4-byte words × 8 sublanes × factor
return lowering_util::FitsInScratchPad(target, sized)
```text

proto 写入后，MLIR `LoopUnrollPass::runOnOperation`（`0x1352ca20`）遍历每个 `scf.for`（`walk<scf::ForOp>`，pre-order）并应用 factor；随后 `VectorUnrollPass` 将宽 vector op 拆分为 native lane width。这里总结这些 MLIR 侧 pass 只是为了闭合数据路径；它们的主体属于 SparseCore lowering 页面。

---

## HLO while-loop 展开

两个开源 pass 运行在 `PostOptimizationPipeline`（`0x1093fd40`）中，这是 latency-hiding scheduler 前的最后一个 HLO pipeline。它们由**独立**开关控制，并且可以同时运行在同一个 module 上。

### `WhileLoopUnroller` — full / double-buffer / auto

由 `*(TpuCompEnv + 4904) != 0` 控制（指向 `xla_while_loop_unroll_count` 的*指针*）。添加时（`AddPass<WhileLoopUnroller, long, bool>` @ `0x1096ee60`），它构建一个 0x30-byte object：

| Offset | Value | 含义 |
|---|---|---|
| `+8` | `*(env + 4904)` | `unroll_count` |
| `+16` | `0` | `wrap_in_trip_count_remainder`（此调用点硬编码 false） |
| `+24` | `64` | `kUnrollTripCountThreshold`（full-unroll trip cap） |
| `+32` | `800` | `kUnrollInstructionCountThreshold`（body-size cap） |
| `+40` | `10000` | `kUnrollExpandFactorThreshold`（trip × instrs expand cap） |

四种模式（DebugOptions string）：`WHILE_LOOP_UNROLLING_NO_UNROLL`（disabled）、`_DOUBLE_BUFFER`（factor 2）、`_FULL_UNROLL`（需要 static trip ≤ 64）、`_AUTO_UNROLL`（仅当 body 包含 collective 时 factor 2）。失败字符串：`"Cannot unroll while loop. The trip count is greater than the threshold:  …  Threshold: "` 和 `"Cannot unroll while loop. Too many instructions in the body: "`。

**`IsLoopUnrollable(HloInstruction*)`（`0x12ee8620`）** — 9 步 legality gate（src `while_loop_unroller.cc`，在命名地址由反编译确认）：

1. opcode == `kWhile`（0x82）— line 1222。
2. 单个 loop-carried tuple（`operands().size() == 1`）— line 1225。
3. 无 control predecessor — line 1238（`"…due to control dependency: "`）。
4. `while_body` 和 `while_condition` 不包含 Send/Recv family `{kSend,kRecv,kSendDone,kRecvDone}` 中的任何 opcode — line 1252（`"…because it contains a send/recv node: "`）。
5. `operand(0)` opcode == `kTuple`（0x81）— line 1259。
6. `while_condition->HasSideEffect() == false` — line 1269。
7. `GetLoopInductionVarTupleIdx` 成功 — line 1277。
8. `HloEvaluator::Evaluate(IV init)` 成功 — line 1287。
9. `MatchTrivialLoopTripCount` 成功 — lines 1295/1299。

全部通过时，它存储 `{while, init_value, trip_count, iv_tuple_idx, is_unrollable=true}`。

> **NOTE — 支撑 exact 4-opcode `flat_hash_set` 的数组没有按字节解码。** IDA 将 rodata adjacency（`unk_AE07CA8` / `unk_AE07CAC`）误标为 ASCII；该 family 由诊断字符串固定为 Send/Recv，而不是通过解码 initializer list 得出。对 family 为 HIGH confidence，对逐字节 opcode list 为 LOW。

### `WhileLoopPipelineUnroller` — 软件流水

由 `EnablePipelinedLoopUnrolling(env)`（`0x1d6b71a0`）控制，它读取 `TpuCompEnv + 752` 处的 `AutoProto`（`xla_tpu_enable_pipelined_loop_unrolling`）：“set” 当且仅当 `(~AutoOr<bool>::FromProtoOrDie(proto) & 0x101) == 0`。在 `PostOptimizationPipeline` 中，它前后分别由 `TpuAnnotateTraceableLoops(true)` 和 `(false)` 包裹。

**`ComputeWhileLoopPipelineDepth(const HloInstruction&)`（`0x12ee0fc0`）** — loop-carry depth = pipeline stage 数 = 某个值在被消费前存活的 iteration 数。它 CHECK `kWhile`（line 44, `"while_instruction.opcode() == HloOpcode::kWhile"`）并 CHECK while-body root 的 *shape* 是 tuple（line 52, `CHECK(while_root->shape().IsTuple())`, `"While Instruction has not been canonicalized to have a tuple shape"`），然后遍历 root tuple 的 operand：读取 `parameter(0)` 且 tuple index `≠ i` 的 `kGetTupleElement`（0x40）是一个 *carry edge*（slot rotation），会记录到 `flat_hash_map<int64,int64>` 中（swiss-table SIMD probe 可见为 `_mm_crc32_u64`/`vpcmpeqb` inner loop）。对 carry-edge graph 做 deque-BFS，并对 chain length 做 binary-GCD reduction，得到 depth。Depth `< 2` ⇒ caller 跳过该 loop。

**`RunImpl`（`0x12ee2200`）** — 转换。对于每个 `depth >= 2` 的 loop，它将 body 克隆成 `depth` 个链式 call stage，并把 trip count 减少 `depth-1`：

```cpp
StatusOr<bool> WhileLoopPipelineUnroller::RunImpl(HloModule* module, threads) {
  for (HloInstruction* loop : while_loops_with_depth_ge_2) {
    int64_t depth = ComputeWhileLoopPipelineDepth(*loop);
    VLOG(1) << "Unrolling: " << loop->name() << " unroll_factor: " << depth;  // :129

    // New outer body "%s.unrolled_%dx": chain `depth` clones of the body as calls.
    HloComputation::Builder b(Format("%s.unrolled_%dx", body->name(), depth));
    HloInstruction* cur = b.AddInstruction(Parameter(0, loop->shape(), "input_tuple"));
    HloComputation* outer = module->AddEmbeddedComputation(b.Build());
    for (int64_t i = 0; i < depth; ++i) {
      HloComputation* stage = module->AddEmbeddedComputation(
          body->Clone(Format("unrolled_%dx_step_%d", depth, i)));
      cur = outer->AddInstruction(Call(loop->shape(), {cur}, stage));
    }
    outer->set_root_instruction(cur);

    HloComputation* new_cond = module->AddEmbeddedComputation(
        cond->Clone(Format("unrolled_%dx", depth)));
    HloInstruction* nw = loop->parent()->AddInstruction(
        While(loop->shape(), new_cond, outer, loop->mutable_operand(0)));

    // Lift depth-1 iterations into the implicit prologue/epilogue.
    Status s = WhileUtil::IncrementWhileLoopTripCount(*nw, /*increment=*/ 1 - depth);  // :176
    nw->set_while_body(outer);
    if (s.ok()) RETURN_IF_ERROR(loop->ReplaceOperandWith(0, nw));
    else        VLOG(1) << "Failed to unroll: " << loop->name();                       // :178
  }
  RETURN_IF_ERROR(FlattenCallGraph().Run(module, threads));   // :188 — inline the stages
  RETURN_IF_ERROR(/*follow-on pass*/.Run(module, threads));   // :190
  return changed;
}

pipelining 语义：body B 在一个 outer body 内变为 depth 个链式 call stage B0→B1→…→B_{depth-1}。因为每个 stage 消费前一 stage 的 output tuple，所以 rotated loop-carry value 会在同一个 outer iteration 中由 stage k 产生，并由 stage k+1 消费——同时有 depth 个原始 iteration in flight。trip count 减少 depth-1（IncrementWhileLoopTripCount(1 - depth)，helper 0x1e3ae7c0）：前 depth-1 次 fill 是隐式 prologue，后 depth-1 次是隐式 drain，都折叠进链式 call 结构中。随后 FlattenCallGraph inline 这些 stage，因此 LLVM modulo scheduler（Bundle Modulo Scheduling）看到并跨 iteration overlap 的 residual counted loop body 大小是 depth × original。

GOTCHA — unroll 与 pipeline 并非互斥。 没有单一的 “unroll-or-pipeline” switch。小 static-trip loop 会被 full/double/auto unrolled（knob A：env+4904）；具有真实 loop-carry depth ≥ 2 的 loop 会额外被软件流水化（knob B：EnablePipelinedLoopUnrolling）。一个 loop 可以被两个 pass 都检查为 candidate；一旦展开为 straight-line code，它就不再匹配 pipeliner 的 kWhile 检查，因此顺序（unroll → pipeline）使二者组合，而不是冲突。

---

Cost-model 交互

在 TPU 路径上，本页的 loop transform 是受约束的，而不是按成本排序的：

• SparseCore window unroll — 贪心：能 fit scratchpad（FitsInScratchPad）的最大 factor。是 resource fit，不是 roofline。
• HLO TileKind — 由 TransferSizeUtil predicate 驱动的正确性/layout 决策，不是成本模型。
• While-loop unrolling — 由 trip-count（≤ 64）/ body-size（≤ 800）/ expand（≤ 10000）阈值控制。
• Pipeline depth — structural property（loop-carry distance），精确计算。

唯一真正成本驱动的 tiling 搜索是 convolution MXU window tiling（MXU cycles + VMEM fit）——并且它不在本页；参见 Dot / Conv → MXU Lowering。

---

示例：6acc60406（“Ironwood”）上的 SparseCore gather loop

给定一个 SparseCore custom-fusion，它在 while (i < 512) 内将 window gather 到 VMEM，目标为 6acc60406（tpu::TpuVersion 5，唯一采用 <5> copy policy 的版本），window slice sizes [1, 8, 128]（BF16，2 B），loop-carry rotation depth 为 3：

1. TileKind（TpuTilingAssignment，post-fusion）：每个 VMEM buffer 的 layout 已经携带 compact tile；fusion 的 kCopy output 得到 GetCopyTileKind → 0x0000（compact in，compact out），若为 SparseCore-laid-out 则为 0x0303。
2. Window unroll（WindowUnrollFactorSelector）：window_elems = 1·8·128 = 1024；bytes = 2048；per-factor size = 8·((2048>>2)+1)·f = 4104·f。选择满足 4104·f ≤ S 的最大 f（例如 S = 64 KiB → f ≤ 15 → f = 8）。BF16（pack=1）的 CustomLoopUnrollPolicy<5> copy factor：MD-vectorizing → 16/1 = 16；否则 inner=128，32·scs_tc（假设 scs_tc=4 → 128）→ 128 % 128 == 0 → 32。selector 将 copy unroll clamp 到 scratchpad 允许的值。
3. 写入 LoopConfig：unrolled_loops += { loop_dim = inner, unroll_factor = f }；vectorizing_shape = [16]。GetNormalizedUnrollFactor：explicit → f；auto（-1）→ bounds[inner] / vectorizing_shape = 128 / 16 = 8。
4. MLIR LoopUnrollPass 按 factor 展开 window 上的 scf.for；remainder = GetLoopPipelineRemainder。
5. HLO pipelining（如果 xla_tpu_enable_pipelined_loop_unrolling）：ComputeWhileLoopPipelineDepth = 3 ⇒ pipeline。Body 被克隆 3× 为 unrolled_3x_step_{0,1,2}，并链式调用；new trip = 512 − 2 = 510；FlattenCallGraph inline 这些 stage。510-iteration counted loop 供给 hardware loop counter；modulo scheduler overlap 这 3 个 stage。

结果： window tile 1×8×128（BF16），copy unroll 32（或被 scratchpad clamp 到 8），scf.for 展开 ×8，while-loop 软件流水 depth 3（trip 510），inner loop 经 modulo-scheduled。

---

手动覆盖标志

Flag	效果
xla_tpu_verify_or_assign_tiling_before_lowering	tri-state：0=skip，1=verify，2=assign（env +0xDFC）；...3/...8 为逐架构变体
xla_tpu_enable_untiled_layout / xla_tpu_untiled_layout_for_1	允许 linear（untiled）layout（TpuTilingRewriter）
xla_tpu_experimental_enable_small_minor_tiling	启用 small-minor special tiling
xla_tpu_enable_large_2nd_minor_layout[_for_x{4,8,16}]	special / large-2nd-minor tiling（post-fusion）
xla_sc_tiles	SparseCore tile count（WindowUnrollFactorSelector 的 long arg）
xla_while_loop_unroll_count	WhileLoopUnroller unroll factor / trip bound（env +4904 pointer；gate A）
xla_tpu_enable_pipelined_loop_unrolling	启用 WhileLoopPipelineUnroller（env +752 AutoProto；gate B）
xla_sc_disable_remainder_loop_pipelining / xla_sc_max_pipelining_stages	SC remainder-loop pipelining disable / stage cap
xla_tpu_scatter_partial_unroll_factor / xla_tpu_unroll_strided_remote_dma	scatter / strided-DMA loop unroll

NOTE — 此二进制文件中的 xla_gpu_* / xla_cpu_* tiling/unroll 标志不在 TPU 路径上。 SymbolicTileAnalysis、TiledHloSchedule、xla_cpu_matmul_tiling_*、xla_gpu_max_kernel_unroll_factor 等随 libtpu.so 一起发布，但只由 xla::cpu::/xla::gpu:: emitter 驱动。重新实现者应在 TPU codegen 中忽略它们。

---

置信度总结

断言	证据
TileKind 是由 GetCopyTileKind pack 的 (input,output) 16-bit pair	反编译 0x13dd0ca0：v18|v23，error lines 66/77，mem-space byte 312，has-bit byte 304
四个 TransferSizeUtil predicate gate compact-vs-SparseCore	在 GetCopyTileKind 中调用；addrs 0x1d6af220/110b7440/1d6af3e0/1d6af2e0
非 kCopy op 通过 GetDefaultLayout 继承 layout；只有 kOutfeed/tuple/async-SC 特殊	RunImpl 0x13dd10a0，subshape visitors 0x13dd26a0/27e0，dtype CHECK :233
VerifyOrAssignTiling tri-state 位于 env +0xDFC（3580）；mode 1→ctor(0)+Verify，mode 2→ctor(1)+Run	反编译 0x10922a20；deepsea_compiler_base.cc:3053/3056
LoopConfig/LoopUnrollConfig 字段 layout（tag、offset、oneof）	wire serializers 0x1d6eade0/0x1d6f2680
Auto unroll factor = loop_bound / vectorizing_shape，带 divisibility CHECK	反编译 0x13d6c1c0；CHECK loop_config_wrapper.cc:358
<3>（viperfish，fallback）SC copy factor 16（elementwise）/ 8（structured）	反编译 0x139173a0：8*elementwise + 8
<5>（6acc60406）SC copy factors（transpose {16,8}；general 32/16；MD 16/pack）	反编译 0x13916fe0 逐字节
tpu::TpuVersion C++ enum = proto−1；dispatch version==5 ⇒ 6acc60406	GetCustomLoopUnrollPolicy 0x13916ec0；TPU_VERSION_* proto descriptor numbers 1..6
Window selector 选择最大的 scratchpad-fitting factor；写入 LoopConfig	Select 0x1385c360，fit test 0x1385c240（8·((bytes>>2)+1)·f）
While-loop unroll thresholds 64/800/10000；gate 位于 env +4904	AddPass 0x1096ee60 object fields；IsLoopUnrollable 0x12ee8620
IsLoopUnrollable 9-step gate；禁止 Send/Recv family	0x12ee8620，src lines 1222–1299
Pipeline depth = loop-carry rotation distance；depth ≥ 2 才 pipeline	ComputeWhileLoopPipelineDepth 0x12ee0fc0
Pipeliner 将 body 克隆成 depth 个链式 call；trip −= depth−1	RunImpl 0x12ee2200；IncrementWhileLoopTripCount(1-depth) 0x1e3ae7c0
Unroll 与 pipeline 独立 gate，可以同时运行	PostOptimizationPipeline 0x1093fd40 中的两个 gate（env+4904，EnablePipelinedLoopUnrolling 0x1d6b71a0）
TpuPostFusionTilingAssignment 传播 special tiling	RunImpl 0x13dd85a0，AcceptsSpecialTiling 0x13dd6580，CanProduceSpecialTiling 0x13dd7760

---

交叉引用

• TPU 编译器 — Part V 导览；这些 loop pass 在 PostOptimizationPipeline（Phase 1）中相对于 layout assignment、fusion 和 scheduler 的位置。
• Fusion Patterns — 在 TpuPostFusionTilingAssignment 重新平铺其 output 之前运行的 fusion；post-fusion stamp 重新平铺的 copy 在这里以及 copy insertion 中引入。
• Dot / Conv → MXU Lowering — 按成本排序的 MXU operand-window tiling（IterateThroughWindowConfigs），本页谨慎避免混淆的“tiling”第二层含义。
• tpu → LLO Lowering — legalizer 会读取本页标记的 TileKind tag，以及 unroller 留下的 loop。
• Tile-Index Expansion — SparseCore 逐 tile memref index algebra（ExpandTiledMemRefs），即 SC tiling 的地址侧。
• Bundle Modulo Scheduling — LLVM 侧 initiation-interval 搜索，对 HLO unroll/pipeline 后留下的 loop 的 LLO bundle stream 做软件流水；“pipelining”的第三层含义。
• 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 / Fusion, dot/conv, tiling — 返回索引