IsaEmitter 注册表
本页中的每个偏移、数值和地址都逐字节精确读取自
libtpu-0.0.40-cp314wheel 中的libtpu.so(BuildID md589edbbe81c5b328a958fe628a9f2207d)。其他版本会有所不同。
摘要
TensorCore 和 BarnaCore 代码生成器不会硬编码哪个 IsaEmitter 子类为给定芯片 lowering 一个 LLO bundle。它们会查询一个进程全局的pair-key 注册表:一个 util_registration::FunctionRegistry,以 std::pair<tpu::TpuVersion, tpu::TpuSequencerType> 为键,存储值是构建 unique_ptr<xla::jellyfish::IsaEmitter> 的工厂。每个 (generation, sequencer) 单元都由一个 google_init_module_*_emitter 翻译单元在 static-init 时注册一次;编译时,IsaEmitterFactory::Create 把两个 enum 打包成一个 64 位键,查找该单元,并调用该单元的工厂来构造 leaf emitter,由它把每个 LLO op lowering 为其 proto bundle。
该注册表是高层级的选择器 —— 它为每个 (gen, seq) 选择 leaf emitter class。随后 leaf 会把每个 LLO opcode 路由到它自己的逐 op EmitX 模板族(proto 填充层),而 per-slot encoders 最后打包位。因此,注册表的两个轴是 TensorCore/BarnaCore 编码路径的主干:第一个轴(TpuVersion)选择硅片代际;第二个轴(TpuSequencerType)选择 bundle 目标所在的芯片 sequencer(TensorCore 或两个 BarnaCore 变体)。总共填充了 12 个单元,分布在 8 个 init modules 中,解析为 6 个不同的 leaf emitter classes。SparseCore sequencers(SparseCoreSequencer / TAC / TEC)出现在 TpuSequencerType enum 中,但没有注册任何单元 —— 它们通过一个独立的 variant-keyed dispatcher 到达自己的 EmitX 模板(见下面的 QUIRK)。
对于重新实现,约定如下:
- 键是一个 64 位打包值:低 dword =
TpuVersion,高 dword =TpuSequencerType—— 在注册侧(每个模块构建该常量)和查找侧(IsaEmitterFactory::Create读取Target+0x398作为 version,并把 sequencer 左移进高 dword)均已确认。 Target+0x398是查找的唯一 version 来源 —— 与其余 target/cost-model 层使用的字段相同。- 注册表存储
shared_ptr<MapValue>;一个MapValue携带工厂闭包(__call_func)和+0x10处的一字节“absent”标志。查找 MISS 是一次硬LogMessageFatal("couldn't create ISA emitter for target: …"),不是软 fallback。 - 12 单元普查:每个
(gen, seq)单元安装一个特定的__call_func包装器,其 lambda 构造一个具体的IsaEmitterleaf。一个包装器可以服务多个代际(Jellyfish/Dragonfish 共享一个 leaf;v6e/v7 共享另一个)。 - v4+
Target类是xla::jellyfish::Target的直接子类(单继承);只有DragonfishTarget : JellyfishTarget是两级链。
| 注册表类型 | util_registration::FunctionRegistry<pair<tpu::TpuVersion, tpu::TpuSequencerType>, unique_ptr<IsaEmitter>(Target const*, CompilerMetadata*, TpuSequencerType, bool, bool, IsaEmitter*, optional<bool>)> |
| 单例 | xla::jellyfish::GetIsaEmitterRegistry @ 0x143f6480(__cxa_guard 保护的函数局部 static;40-B 对象零初始化) |
| 注册 | FunctionRegistry::Register @ 0x140c2360(mutex 锁保护的 flat_hash_map 插入) |
| 查找 | FunctionRegistry::Get @ 0x140af4e0;raw_hash_set find @ 0x140af5e0 中的键比较 |
| 工厂入口 | xla::jellyfish::IsaEmitterFactory::Create @ 0x140af220 |
| 键布局 | uint64 = (uint32 TpuSequencerType << 32) | (uint32 TpuVersion) —— [+0]=version,[+4]=seqtype |
| Version 来源 | Target+0x398(在 Create+0x6f 读取) |
| 已填充单元 | 12(8 init modules → 6 leaf classes) |
| Map 策略 | FlatHashMapPolicy<pair<TpuVersion,TpuSequencerType>, shared_ptr<MapValue>> |
| MISS | LogMessageFatal "couldn't create ISA emitter for target:"(Create+0x259) |
| 置信度 | CONFIRMED(反编译验证),除非行内另有说明 |
两个轴
该注册表是一个二维表:(TpuVersion) × (TpuSequencerType) → IsaEmitter leaf。两个轴都不是自由索引 —— 二者都是由硅片定义的 enum,只有当该组合是实际硬件上的真实引擎时,单元才存在。
轴 1 —— TpuVersion 是硅片代际,与 codec metadata 和 cost model 使用的是同一个六值 enum。tpu::TpuVersionToString(0x20b3a480)索引一个 6 项指针表,并对任何 ≥ 6 的 ordinal trap:
TpuVersion | Codename | Public name |
|---|---|---|
| 0 | jellyfish | TPU v2 |
| 1 | dragonfish | TPU v3 |
| 2 | pufferfish | TPU v4 |
| 3 | viperfish | TPU v5p (+ v5e lite) |
| 4 | ghostlite | TPU v6e |
| 5 | 6acc60406 | TPU v7 |
codename 字符串直接读取自 tpu::TpuVersionToString 表(off_22011BF0):ordinal 0..5 解析为 jellyfish、dragonfish、pufferfish、viperfish、ghostlite、6acc60406。public-name 列遵循 每代比较矩阵 中规范的 codename → marketing-name 映射;6acc60406 是唯一一个二进制中没有携带 public-name 字符串的代际(字面量 Trillium/Ironwood 在 libtpu.so 中完全不存在 —— 6acc60406 是该代际唯一的 codename)。
轴 2 —— TpuSequencerType 表示 bundle 目标是芯片中的哪个 sequencer。tpu::TpuSequencerTypeToString(0x20b362e0)是一条单指令 —— return off_22010DE0[ordinal] —— 一个按 ordinal 索引、含八个条目的指针表:
TpuSequencerType | 名称 | 是否注册单元? |
|---|---|---|
| 0 | TensorCoreSequencer | 是(每代) |
| 1 | BarnaCoreSequencer | 是(v0/v1/v2) |
| 2 | BarnaCoreAddressHandler | 是(v0/v1/v2) |
| 3 | SparseCoreSequencer | 否(独立路径) |
| 4 | SparseCoreTileAccessCoreSequencer (TAC) | 否 |
| 5 | SparseCoreTileExecuteCoreSequencer (TEC) | 否 |
| 6 | IMEM | 否 |
| 7 | VIMEM | 否 |
笛卡尔积是稀疏的。只有命名该代际上实际存在 sequencer 的 (gen, seq) 对会被注册:v0/v1/v2 带有一个 TensorCore 加 BarnaCore 引擎;v3/v4/v5 在这条路径上只有 TensorCore(它们的 SparseCore lowering 是一个独立 dispatcher)。因此 6 × 8 = 48 个网格位置中只有 12 个活动单元。
注意 —— 两个 enum 在二进制中是独立的,但在 key word 中耦合。version 来自
Target+0x398(target 对象的运行时属性);sequencer type 是调用者为正在 lowering 的引擎显式传入的参数。重新实现如果从 version 推导 sequencer,或反过来推导,会错误地构造查找键。
键布局和查找路径
IsaEmitterFactory::Create(0x140af220)是唯一入口。它构建键、查找键,然后要么调用该单元的工厂,要么终止。反编译后的函数体精确展示了打包和 MISS 检查:
// xla::jellyfish::IsaEmitterFactory::Create @ 0x140af220 (decompiled, trimmed)
__int64 Create(const Target *target, CompilerMetadata *md,
TpuSequencerType seq, bool compact, ...) {
if (seq == 2 /* BarnaCoreAddressHandler */) {
// BarnaCoreAddressHandler rejects parallel codegen and compact emit (Fatal)
...
}
Registry *reg = GetIsaEmitterRegistry(); // the singleton
// KEY PACK: low dword = Target+0x398 (version); high dword = seq
uint64_t key = *(uint32_t*)((char*)target + 0x398)
| ((uint64_t)seq << 32);
MapValue *v = FunctionRegistry::Get(reg, &key); // raw_hash_set find @ 0x140af5e0
if (v->absent /* byte @ MapValue+0x10 == 1 */)
LogMessageFatal("couldn't create ISA emitter for target: ", target->name);
// INVOKE the cell's factory __call_func (slot v[2]) -> constructs the leaf
return v->call_func(v, target, md, seq, compact, ..., key_lo16);
}
```text
该函数体中的三个事实都有加载锚点:
- **`Target+0x398` 是 version 来源。** 反编译器将其渲染为 `*((unsigned int *)v10 + 230)` —— `230 × 4 = 0x398`。这与 [target/cost-model 层](../targets/tpuhal-class-hierarchy.md) 作为代际选择器读取的 `Target+0x398` 相同,因此 IsaEmitter 注册表和 cost model 都基于*同一个*字段构建键。
- **MISS 是 fatal。** `Get` 返回的 `MapValue` 在 `+0x10` 处携带一个“absent”哨兵字节;当该字节置位时,`Create` 触发 `LogMessageFatal`。不存在默认 emitter —— 未注册的 `(gen, seq)` 会中止编译。
- **Sequencer 2 会被提前特判。** 在查找前,`seq == BarnaCoreAddressHandler` 会用自己的 fatal 拒绝 parallel codegen 和 compact emit —— 这是注册表值本身无法表达的逐引擎约束。
`raw_hash_set` find(`0x140af5e0`)中的键比较会分别比较低 dword(version)和 `+4` dword(sequencer)—— 这直接确认 64 位键是两个打包的 int32,version 在低半部。
---
## 注册侧
每个单元都由一个调用 `FunctionRegistry::Register`(`0x140c2360`)的 `google_init_module_*_emitter` 函数在 static-init 时安装一次。`Register` 获取一个 `mutex` 锁,在堆上分配一个 0x48 字节的 `MapValue`(携带工厂闭包),并把它插入以 `pair` 为键的 `flat_hash_map`;重复键会触发 `LogMessageFatal` "Registration failed; key already exists in registry"。模块构建的键常量与查找时重建的同一个 packed `uint64` 相同。
`jellyfish_emitter` 模块是最清晰的例子:它连续注册四个单元,全部安装**同一个** `JellyfishEmitter` 工厂闭包(`$_0`),只是键常量不同:
```c
// google_init_module_jellyfish_emitter @ 0x213ecdc0 (decompiled, trimmed)
key = 0; Register(GetIsaEmitterRegistry(), &key, JellyfishEmitter_$_0, ...); // (v0, seq0)
key = 0x100000000; Register(GetIsaEmitterRegistry(), &key, JellyfishEmitter_$_0, ...); // (v0, seq1)
key = 1; Register(GetIsaEmitterRegistry(), &key, JellyfishEmitter_$_0, ...); // (v1, seq0)
key = 0x100000001; Register(GetIsaEmitterRegistry(), &key, JellyfishEmitter_$_0, ...); // (v1, seq1)key = 0x100000000 是 (version 0, seqtype 1);key = 1 是 (version 1, seqtype 0)。高 dword 是 sequencer,低 dword 是 version —— 与 Create 打包的布局完全一致。同一种写法出现在每个模块中;每代模块在自己的 Register 调用前都会加载一个单独的键常量(Viperfish 为 mov [rbp-8], 3,Ghostlite 为 4,6acc60406 为 5)。
注意 —— 注册的闭包是一个工厂,不是 leaf 本身。
MapValue存储围绕 lambda 的__call_func包装器;该 lambda 在Create内按需构造具体 leaf。因此一个闭包可以服务多个键 —— 上面的四个 Jellyfish/Dragonfish 单元全都共享一个JellyfishEmitter闭包,所以一个 leaf class 同时 lower v0 和 v1 的 TensorCore 与 BarnaCore sequencer。
12 单元普查
全节扫描发现正好十二个馈入该注册表的 Register 调用点,分布在八个 init modules 中。每个单元的 leaf 是该单元 __call_func lambda 构造的 IsaEmitter 子类。
| # | key (u64) | (TpuVersion, TpuSequencerType) | Init module | Leaf emitter |
|---|---|---|---|---|
| 1 | 0x000000000 | (0 jellyfish, 0 TensorCore) | jellyfish_emitter | JellyfishEmitter |
| 2 | 0x100000000 | (0 jellyfish, 1 BarnaCoreSequencer) | jellyfish_emitter | JellyfishEmitter |
| 3 | 0x200000000 | (0 jellyfish, 2 BarnaCoreAddressHandler) | barna_core_address_handler_emitter | BarnaCoreAddressHandlerEmitter |
| 4 | 0x000000001 | (1 dragonfish, 0 TensorCore) | jellyfish_emitter | JellyfishEmitter |
| 5 | 0x100000001 | (1 dragonfish, 1 BarnaCoreSequencer) | jellyfish_emitter | JellyfishEmitter |
| 6 | 0x200000001 | (1 dragonfish, 2 BarnaCoreAddressHandler) | barna_core_address_handler_emitter | BarnaCoreAddressHandlerEmitter |
| 7 | 0x000000002 | (2 pufferfish, 0 TensorCore) | pufferfish_tensorcore_emitter | PufferfishTensorCoreEmitter |
| 8 | 0x100000002 | (2 pufferfish, 1 BarnaCoreSequencer) | pufferfish_barnacore_sequencer_emitter | PufferfishBarnaCoreSequencerEmitter |
| 9 | 0x200000002 | (2 pufferfish, 2 BarnaCoreAddressHandler) | pufferfish_barnacore_channel_emitter | PufferfishBarnaCoreChannelEmitter |
| 10 | 0x000000003 | (3 viperfish, 0 TensorCore) | viperfish_tensorcore_emitter | ViperfishTensorCoreEmitter |
| 11 | 0x000000004 | (4 ghostlite, 0 TensorCore) | ghostlite_tensorcore_emitter | GhostliteTensorCoreEmitter |
| 12 | 0x000000005 | (5 6acc60406, 0 TensorCore) | 6acc60406_tensorcore_emitter (sub_213ED1C0) | GhostliteTensorCoreEmitter(复用) |
该表的形状反映了硅片事实:
- v0/v1(Jellyfish/Dragonfish) 各有三个单元 —— 一个 TensorCore 加芯片的两个 BarnaCore sequencer 角色 —— 由两个 leaf classes 服务(
JellyfishEmitter服务 TC + BarnaCoreSequencer;BarnaCoreAddressHandlerEmitter服务 address-handler)。barna_core_address_handler_emitter模块(0x213ed040)在一个函数体中以两个显式键常量安装单元 3 和 6 —— 先0x200000000,再0x200000001,二者都使用同一个$_0闭包 —— 因此两个 address-handler 单元是字节确认的,而不是推断的。 - v2(Pufferfish) 也有三个单元,但其 BarnaCore 被拆分为一个 sequencer emitter 和一个 channel emitter(单元 8 和 9),二者各自是不同的 leaf。
- v3/v4/v5(Viperfish/Ghostlite/6acc60406) 只有一个 TensorCore 单元。v5p+ 上没有 BarnaCore;它们的 SparseCore 走独立路径。
- 单元 12 复用单元 11 的 leaf。 6acc60406(TPU v7)TensorCore 单元安装与 Ghostlite(TPU v6e)相同的
GhostliteTensorCoreEmitter工厂 —— 这与 v5-ordinal 代际复用GhostliteTarget的运行时事实一致。代际合并发生在 leaf-class 层;gfc-vs-glc 编码器拆分发生在 codec 内部的下游。
单元 11 和 12 从两个不同的翻译单元注册,而不是同一个。具名的 ghostlite_tensorcore_emitter init 函数(0x213ed160)只注册单元 11(key 4)然后返回;单元 12(key 5)由相邻的 static-init 函数(sub_213ED1C0)注册,该函数编译自一个独立的 6acc60406_tensorcore_emitter 源文件,并安装同一个 GhostliteTensorCoreEmitter 工厂。两个 __call_func thunk(单元 12 为 sub_14398B60,单元 11 为 0x142a04c0)都会 operator new(0x2A0u),随后调用同一个 GhostliteTensorCoreEmitter::GhostliteTensorCoreEmitter ctor —— 两个单元共享 leaf class,但来自两个 TU。
为什么只存在 TensorCore/BarnaCore 单元 {#why-only-tensorcorebarnacore-cells-exist}
TpuSequencerType enum 暴露了三个 SparseCore sequencers(SparseCoreSequencer = 3,TAC = 4,TEC = 5)和两个 memory sequencers(IMEM = 6,VIMEM = 7),但没有模块为其中任何一个注册单元。SparseCore 引擎通过完全独立的 dispatcher 到达自己的 EmitX 模板。
QUIRK —— SparseCore 引擎完全绕过 pair-key 注册表。 SparseCore(SCS/TAC/TEC)lowering 由
xla::tpu::sparse_core::code_generator驱动,这是一条两级路径,基于 chip-parts variant(而不是Target+0x398)构建键:RunCodeGen→ gen-switchMakeTpuCoreProgram→ per-gen 模板实例化(MakeTpuCoreProgram<{Viperfish,Ghostlite}Emitter, …>)→Emitter::ConsumeProgram→ 基于 MCInst opcode 的每个 bundle 类型Consume*Instructionjump-table →EmitX<Bundle, Op>。逐指令引擎(SCS/TAC/TEC)由 section-name classifier 选择,而不是由注册表Get选择。重新实现如果预期 pair-key 注册表拥有所有 sequencer type,会发现 SparseCore 半边缺失,必须单独建模 variant-keyed dispatcher。
实际分工是:pair-key 注册表拥有编码路径中的 TensorCore + BarnaCore 半边(上面的 12 个单元);variant-keyed code_generator 拥有 SparseCore 半边。二者最终都到达同一个 per-slot encoder → BitCopy 位打包阶段。
Target 基类链
version 轴由一个 Target 类层次支撑:IsaEmitterFactory::Create 读取 Target+0x398 来为注册表构建键,而每代 Target 子类携带该字段以及 emitter 要查询的 chip-parts profile。该层次全程为单继承(每个 typeinfo 都是 __si_class_type_info),根是抽象的 xla::jellyfish::Target:
xla::jellyfish::Target (abstract root, __class_type_info)
├── JellyfishTarget (ordinal 0)
│ └── DragonfishTarget (ordinal 1) ← the ONLY two-level chain
├── PufferfishTarget (ordinal 2) ← direct
├── ViperfishTarget (ordinal 3) ← direct
└── GhostliteTarget (ordinal 4; reused by ordinal 5 / 6acc60406, no separate class)
xla::jellyfish::SparseCoreTarget (parallel abstract root)
├── ViperfishSparseCoreTarget
└── GhostLiteSparseCoreTarget
```text
| Class | ZTI @ | base class | object size | ordinal / public |
|---|---|---|---|---|
| `Target` (root) | `0x21ccef00` | — | — | — |
| `JellyfishTarget` | `0x21cc7420` | `Target` | `0x958` (2392 B) | 0 / v2 |
| `DragonfishTarget` | `0x21cc6ba8` | `JellyfishTarget` | `0x958` | 1 / v3 |
| `PufferfishTarget` | `0x21cc7d38` | `Target` | `0x950` (2384 B) | 2 / v4 |
| `ViperfishTarget` | `0x21cc8f78` | `Target` | `0x950` | 3 / v5p |
| `GhostliteTarget` | `0x21cc85f8` | `Target` | `0x950` | 4 / v6e (+5/v7 reuse) |
| `SparseCoreTarget` (root) | `0x21ccef10` | — | — | — |
| `ViperfishSparseCoreTarget` | `0x21cc9080` | `SparseCoreTarget` | — | v5p SC |
| `GhostLiteSparseCoreTarget` | `0x21cc8700` | `SparseCoreTarget` | — | v6e SC |
8 字节对象大小差异(v0/v1 为 `0x958`,v2+ 为 `0x950`)来自 `JellyfishTarget` 额外的 `+0x950` `Performance*` 字段(其 ctor 中构建的 Jf cost-model 对象);v4+ 通过不同路径构建自己的 performance 对象,并省略该字段。
每代 ctor 在 base `Target::Target` 调用和 vtable patch 之外只贡献少量 `this` 写入。`PufferfishTarget::PufferfishTarget`(`0x1d493840`)具有代表性,并且已由反编译确认:
```c
// PufferfishTarget::PufferfishTarget @ 0x1d493840 (decompiled)
Target::Target(this, version, variant_name, ..., CreateDefaultTargetEnv(chip_parts));
*(void**)this = off_21CC74E8; // vtable patch -> PufferfishTarget
*((uint32_t*)this + 0x14F) = 5; // +0x53c : chip-generation code = 5
*((uint32_t*)this + 0x245) = 1; // +0x914 : config word = 1
new_divisor = ConstantDivisor(16); // +0x938 : lane/tile divisor = 16
*((void**)this + 0x127) = new_divisor;单个 Target::Target base 调用确认 PufferfishTarget 直接派生自 Target —— 不存在逐架构中间层。只有 v0→v1 这一对成链(Dragonfish 复用 Jellyfish 的 ctor 并 patch vtable)。每代 ctor 值就是硅片 profile:+0x53c 是 chip-generation code(Pufferfish 为 5,Viperfish/Ghostlite 为 7),+0x914 是一个 config word,+0x938 是一个 util::math::ConstantDivisor,其 divisor 是每代的 lane/tile 数量(Jellyfish / Pufferfish / Viperfish+ 分别为 8 / 16 / 32)。
注意 —— 每代 ctor 不是设置大部分 target 字段的位置。共享的
Target::Init(0x1d60fc20,所有代际共用)会写入+0x3cc..+0x948范围内 173 个不同 target 字段;每代的值来自Init内部分派的 chip-parts variant-visit lambda,而不是来自不同的Init函数体。上面的 ctor 只盖上少数结构上按 class 区分的字段。
注册表如何嵌入编码流水线
注册表是 LLO bundle 与其打包字节之间三层中的第一层:
Target+0x398 (TpuVersion) + TpuSequencerType (arg)
│
▼
1. IsaEmitterFactory::Create ──pack key──▶ FunctionRegistry::Get ──▶ IsaEmitter leaf
│ (per (gen,seq))
▼
2. leaf::Emit<op> ──▶ EmitX<Bundle,Op> family (proto submessage population)
│
▼
3. <Slot>Encoder::Encode ──▶ BitCopy(buf, abs_bit, &field, 0, width) (bit packing)
```text
第 1 层(本页)选择哪个 leaf emitter 处理一个 `(gen, seq)`。第 2 层是 leaf 的逐 op `EmitX` 模板族,它会为每个 op 填充一个 proto bundle 子消息 —— 即 [proto 填充层](v5plus-emitx-bit-positions.md)。第 3 层是 per-slot encoder,最后通过 `BitCopy` 把绝对位位置写入 [固定宽度 bundle word](bundle-model-overview.md)。注册表本身不贡献任何位;它是纯 dispatch。它的角色与 LLVM-MC 侧的表选择完全类似 —— 但 [MC emitter](mc-emitter.md) 通过 jump table 基于 opcode 选择,而 IsaEmitter 注册表通过 hash map 基于 `(gen, seq)` 选择,两条路径互补:MC emitter 对每个 TensorCore/V5+ opcode 返回全零,正是因为它们的真实字节来自本注册表选择进入的 proto-bundle 路径。
---
## 交叉引用
- [V5+ EmitX 位位置](v5plus-emitx-bit-positions.md) —— leaf emitter 供给的逐 op `EmitX` → `<Slot>Encoder::Encode` → `BitCopy` 位打包阶段。
- [MC-Emitter](mc-emitter.md) —— `getBinaryCodeForInstr`;这个注册表的 leaf 编码的每个 opcode 在互补的 LLVM-MC 路径中都会返回全零。
- [Bundle 模型](bundle-model-overview.md) —— 编码路径将 slot 铺入其中的每代固定宽度 bundle word。
- [TpuHal 类层次](../targets/tpuhal-class-hierarchy.md) —— 所有 target/codec/cost-model 树都分派依据的 `tpu::TpuVersion` 轴,包括本注册表用于构建键的 `Target+0x398` 字段。
- [每代比较矩阵](../appendix/per-gen-comparison-matrix.md) —— 本页 version 轴所解析到的规范 codename ↔ `TpuVersion` ordinal ↔ public-name 映射(jellyfish v2 … 6acc60406 v7)。