ArchRegno 编号
本页中的每个偏移、值和地址都按字节精确读取自
libtpu-0.0.40-cp314wheel 中的libtpu.so(BuildID md589edbbe81c5b328a958fe628a9f2207d)。其他版本会有所不同。
摘要
TensorCore 编码器在发出一个 bundle slot 时,需要为每个架构寄存器取得一个单一的稠密整数,但 LLO IR 携带的是由 RegisterType(preg / sreg / vmreg / vreg)加该类别内 regno 描述的虚拟寄存器。桥接层是 RegisterNumbering:一个按 Target 创建的对象,为每个 (RegisterType, regno) 分配一个连续的架构寄存器编号(arch_regno),同时维护正向映射((type, regno) → arch_regno,由 ToRegNum 读取)和反向向量(arch_regno → (type, regno),由 ToArchRegno 读取)。由于不同硅片代际的寄存器文件大小不同,编号会在每次构造 Target 时根据四个代际特定的计数字段重新构建。本页记录这个构建过程,即 Target::InitRegisterNumbering 和 RegisterNumbering::Init / AddRegister,以及塑造它的各代计数输入。
如果你熟悉 LLVM,最接近的类比是一个分配 MCRegister 编号的 TargetRegisterInfo,但这里的布局不是由 TableGen 生成,而是在运行时把四个按类别的寄存器计数做前缀和,按固定类别顺序压入一个连续的架构寄存器空间,其中 arch_regno 0 保留为空哨兵。每个 Target 有三个这样的 RegisterNumbering 对象,每个 sequencer 一个,并且各自使用计数字段的不同子集来确定大小。这个编号之下的物理层,即 ArchRegisterInstance 解析的 50 成员 ArchRegister 枚举,记录在 ResultFifo 和 ArchRegister 枚举;本页描述的是其上一层。
对于重新实现,契约是:
- 每个
Target三个RegisterNumbering对象的布局、每个对象消耗的四个计数来源字段,以及从TargetEnvironment取得的按类别名称字符串。 RegisterNumbering::Init:把四个类别计数前缀和到total_registers_,把arch_regno 0保留为kNoRegister,然后按kPreg → kSreg → kVmreg → kVreg顺序分配连续的arch_regno,同时构建两种映射。ToArchRegno/ToRegNum/ToArchRegString:读取路径、边界,以及(type, regno)的打印格式。- 按 opcode 的 variant 分类(
LloOpcodeUsesTranspose/UsesRpu/IsRpuControl/IsRpuResult),以及每个 op 的ArchRegister读/写集合如何馈入CrossXluOperationsDataDependencyTracker。
| 构建入口 | Target::InitRegisterNumbering @ 0x1d614200 |
| 按类别编号器 | RegisterNumbering::Init @ 0x1d622520 |
| 寄存器分配器 | RegisterNumbering::AddRegister @ 0x1d622bc0 |
| 反向读取 | RegisterNumbering::ToArchRegno @ 0x1275f580 ([this+0x80][idx]) |
| 正向读取 | RegisterNumbering::ToRegNum @ 0x1d5a9000 ([this+0x98] map) |
| 名称打印 | RegisterNumbering::ToArchRegString @ 0x1275e2a0 |
| 类别顺序 | kAllocatableRegisterTypes @ 0xadf7f54 = {kPreg=1, kSreg=2, kVmreg=3, kVreg=4} |
RegisterNumbering 大小 | 0x130 字节 |
| 源文件 | platforms/xla/service/jellyfish/register_numbering.{h,cc} |
Target::InitRegisterNumbering —— 构建过程
目的
InitRegisterNumbering @ 0x1d614200 会在每个 Target 上调用一次,用来填充嵌入在 Target 中的三个 RegisterNumbering 对象。每个对象捕获一个 sequencer 的寄存器文件。对象偏移和步长均为字节精确:
| 映射 | RegisterNumbering this | Sequencer 角色 | 存在的类别 |
|---|---|---|---|
| 1 | Target + 0x008 | Primary / TensorCore | 全部四类(preg, sreg, vmreg, vreg) |
| 2 | Target + 0x138 | 非 primary(仅 preg+sreg) | preg, sreg |
| 3 | Target + 0x268 | 非 primary(仅 vmreg+vreg) | vmreg, vreg |
步长 0x130 就是 RegisterNumbering 大小(0x138 - 0x008 = 0x268 - 0x138 = 0x130)。反编译确认了在 this+8、this+312(0x138)和 this+616(0x268)处的三次 RegisterNumbering::Init 调用。
算法
每个映射都是一个新构建的 absl::flat_hash_map<RegisterType, pair<int count, string name>>,含四个条目(键 1..4),传给 RegisterNumbering::Init。按类别的 count 来自 Target 寄存器计数字段块;按类别的 name 来自 TargetEnvironment。
function InitRegisterNumbering(target): // sub_1d614200
env = target[0x940] // TargetEnvironment
// four per-class name strings, copied from the environment:
name_preg = env[0x18]; name_sreg = env[0x20]
name_vmreg = env[0x28]; name_vreg = env[0x30]
// ---- Map 1: primary / TensorCore sequencer (all 4 classes) ----
map = { kPreg: (target[0x4a4], name_preg), // *((DWORD*)target+297)
kSreg: (target[0x498], name_sreg), // *((DWORD*)target+294)
kVmreg: (target[0x4a0], name_vmreg), // *((DWORD*)target+296)
kVreg: (target[0x49c], name_vreg) } // *((DWORD*)target+295)
RegisterNumbering::Init(target + 0x008, map)
// ---- Map 2: preg+sreg-only sequencer ----
map = { kPreg: (target[0x4c0], ...), kSreg: (target[0x4b4], ...),
kVmreg: (0, ...), kVreg: (0, ...) }
RegisterNumbering::Init(target + 0x138, map)
// ---- Map 3: vmreg+vreg-only sequencer ----
map = { kPreg: (0, ...), kSreg: (0, ...),
kVmreg: (target[0x4bc], ...), kVreg: (target[0x4b8], ...) }
RegisterNumbering::Init(target + 0x268, map)这些计数字段偏移通过反编译中的 *((_DWORD *)this + N) 索引算术得到字节级确认:this+294 = 0x498(Sreg)、+295 = 0x49c(Vreg)、+296 = 0x4a0(Vmreg)、+297 = 0x4a4(Preg);非 primary 映射使用 +301 = 0x4b4、+302 = 0x4b8、+303 = 0x4bc、+304 = 0x4c0。整个 [0x498..0x4c0] 块由 Target 构造函数清零,并在 Target::Init 期间从 chip-parts sequencer 描述符填充。
函数映射
| 函数 | 地址 | 角色 |
|---|---|---|
Target::InitRegisterNumbering | 0x1d614200 | 构建 3 个 RegisterNumbering 对象 |
Target::RegisterCount | 0x1d617120 | 同一组 4 个计数字段的读取侧,按 (seq, type) 取键 |
Target::SregCount | 0x1d6152c0 | mov 0x498(rdi),eax —— 确认 Sreg 偏移 |
Target::VregCount | 0x1d6152e0 | mov 0x49c(rdi),eax —— 确认 Vreg 偏移 |
Target::RegisterCount @ 0x1d617120 是读取侧访问器,并逐字节交叉确认了 (seq, type) → offset 映射:sequencer 0 映射 {kPreg→0x4a4, kSreg→0x498, kVmreg→0x4a0, kVreg→0x49c};sequencer 1 为 {kPreg→0x4c0, kSreg→0x4b4};sequencer 2 为 {kVmreg→0x4bc, kVreg→0x4b8},与三个映射完全一致,因此每个 RegisterNumbering 对象都精确对应一个 sequencer 的寄存器文件。
注意 —— 按类别的名称字符串可由每个
TargetEnvironment配置(env+0x18/0x20/0x28/0x30),默认是RegisterTypeToMnemonic的{p, s, vm, v}助记符。重新实现必须从环境读取名称,而不是硬编码,尽管所有已发布 target 都使用默认值。陷阱 —— 第二和第三个
RegisterNumbering对象是部分对象:映射 2 只有 preg+sreg 计数,映射 3 只有 vmreg+vreg,另外两个类别的计数设为 0。因此非 primary sequencer 的RegisterNumbering只为寄存器类别的严格子集编号。这两个对象的精确TpuSequencerType标签(BarnaCore address-handler 还是 SparseCore-tile sequencer)是从计数子集和RegisterCount的 sequencer 分支推断出来的,而不是来自已解码的 sequencer-type 表。
RegisterNumbering::Init —— 为类别编号
目的
RegisterNumbering::Init @ 0x1d622520 接收一个 sequencer 的 flat_hash_map<RegisterType, (count, name)>,并为每个寄存器分配连续的 arch_regno,同时构建映射的两个方向。
对象布局
struct RegisterNumbering { // sizeof 0x130
int32 total_registers_; // +0x00 prefix-sum of the 4 class counts
InlinedBitVector<128> allocatable[5]; // +0x08 per-RegisterType masks, stride 0x18
vector<pair<RegisterType,int>> idx_to_regno_; // +0x80 arch_regno -> (type, regno)
flat_hash_map<pair<RegisterType,int>, int> regno_to_idx_; // +0x98 (type, regno) -> arch_regno
};按类型的 bit-vector 子结构从 +0x08 开始,步长为 0x18(AddRegister 计算 &this[8 + 0x18*type]);五个类型占用 0x08..0x80。反向向量位于 +0x80(int* 索引 32),正向 map 位于 +0x98(int* 索引 38),二者都由 AddRegister 的 push_back 和 insert 目标逐字节确认。
算法
function RegisterNumbering::Init(this, class_count_map): // sub_1d622520
this.total_registers_ = 1
running = 1
for type in kAllocatableRegisterTypes: // {kPreg, kSreg, kVmreg, kVreg}, stride 4, bound 16
running += class_count_map.at(type).count // map "at" lookup; .count = pair.first
this.total_registers_ = running // [this+0] = prefix-sum total
// NOTE: total starts at 1 — arch_regno 0 is reserved below, so the count
// includes the null sentinel.
AddRegister(this, kNone=0, regno=0, arch_regno=0, is_pseudo=false) // reserves arch_regno 0
next_arch_regno = 1
for type in kAllocatableRegisterTypes: // SECOND pass, same order
mnemonic = RegisterTypeToMnemonic(type)
spec = RangeSpec(...) // config-driven name/range filter
for k in [0, class_count_map.at(type).count):
name = StrCat(mnemonic, FastIntToBuffer(k)) // "v3", "s12", ...
is_pseudo = RangeSpec::Match(spec, k, name, 1)
AddRegister(this, type, regno=k, arch_regno=next_arch_regno + k, is_pseudo)
next_arch_regno += class_count_map.at(type).count
// post-conditions (assertion strings in the binary):
assert index_to_regno_[kNoRegister].first == RegisterType::kNone
assert GetMask(RegisterType::kNone).count() == 0
assert total_registers_ == reg_num // reg_num = final next_arch_regno分配顺序由 kAllocatableRegisterTypes @ 0xadf7f54 固定(四个 int32 = {1, 2, 3, 4}):架构寄存器空间布局为 [sentinel] [Preg block] [Sreg block] [Vmreg block] [Vreg block],每个块的大小由对应类别计数决定。
RegisterNumbering::AddRegister
function AddRegister(this, type, regno, arch_regno, is_pseudo): // sub_1d622bc0
assert type <= 4 // "reg_class < kNumberRegisterTypes"
assert arch_regno < this.total_registers_ // "reg_num < total_registers_"
// 1. mark this regno allocatable in the per-type bit vector at &this[8 + 0x18*type]
InlinedBitVector::resize(&this.allocatable[type], total_registers_)
// 2. reverse map: push (type, regno) at index arch_regno
this.idx_to_regno_.push_back({type, regno}) // [this+0x80]
// 3. forward map: (regno<<32 | type) -> arch_regno
this.regno_to_idx_.insert((regno << 32) | type, arch_regno) // [this+0x98]
// 4. if not the sentinel and not a pseudo, record it in the per-type allocatable list
if type != kNone and not is_pseudo:
set_bit(per_type_mask[type], arch_regno)
per_type_allocatable[type].push_back(arch_regno)函数映射
| 函数 | 地址 | 角色 |
|---|---|---|
RegisterNumbering::Init | 0x1d622520 | 前缀和 + 双遍分配 |
RegisterNumbering::AddRegister | 0x1d622bc0 | 分配一个 arch_regno,构建两种映射 |
RegisterTypeToMnemonic | 0x1d640600 | 名称的类别前缀(p/s/vm/v) |
FastIntToBuffer | 0x211719e0 | 为 regno 后缀做整数→ASCII |
RangeSpec::Match | 0x1d624c80 | 配置驱动的按 regno 过滤器(设置 is_pseudo) |
~RegisterNumbering | 0x1d491e60 | 释放 bit-vector / map 缓冲区(确认 0x130 大小) |
特性 —— 前缀和从
total_registers_为 1 开始,第二遍从 1 开始分配arch_regno,而AddRegister(kNone, 0, 0, false)占用arch_regno 0。因此索引 0 总是kNoRegister哨兵,最终分配计数等于total_registers_,断言total_registers_ == reg_num会强制这一点。从 0 开始编号的重新实现会与哨兵冲突并失败此检查。注意 ——
RangeSpec::Match控制哪些 regno 会被编号,以及每个是否为 pseudo(从可分配 mask/list 中排除,但仍会编号)。它的配置来源是TargetEnvironment的寄存器名称/范围过滤器;该机制返回传给AddRegister的is_pseudo标志。因此名义类别计数可能大于可用(可分配)计数。
按代际的寄存器文件大小
编号所消耗的四个计数字段由每个 TpuVersion 的 chip-parts sequencer 描述符填充。字段偏移和前缀和布局是字节精确的;下面的数值计数是按代际取键的 chip-parts 值。
| 代际(代号) | TpuVersion | SREG (0x498) | VREG (0x49c) | VMREG (0x4a0) | PREG (0x4a4) | Arch regs(近似) |
|---|---|---|---|---|---|---|
| v2 / v3 / v4 (jellyfish / dragonfish / pufferfish) | kJellyfish=0, kDragonfish=1, kPufferfish=2 | 32 | 32 | 8 | 15 | ~88 (+1 sentinel) |
v5p (+v5e lite) / v6e / v7 (viperfish / ghostlite / 6acc60406) | kViperfish=3, kGhostlite=4, gen-5=5 | 32 | 64 | 16 | 14 | ~127 (+1 sentinel) |
TpuVersion 枚举有六个值:{kJellyfish=0, kDragonfish=1, kPufferfish=2, kViperfish=3, kGhostlite=4, <gen-5>=5}(第六个,即 v7 代际,在此二进制中代号为 6acc60406)。架构寄存器总数是 1 + PREG + SREG + VMREG + VREG(+1 是 arch_regno 0)。分配顺序始终是 Preg → Sreg → Vmreg → Vreg,因此 Preg 块占据 arch_regno [1, 1+PREG),Sreg 块占据 [1+PREG, 1+PREG+SREG),依此类推。
计数字段偏移(0x498/0x49c/0x4a0/0x4a4) | SregCount/VregCount + RegisterCount 跳转表解码 |
前缀和总数,Preg→Sreg→Vmreg→Vreg 块顺序 | Init 反编译 |
| 每个代际的数值计数 | Chip-parts sequencer 描述符 |
特性 —— VREG 从 32(v2/v3/v4)翻倍到 64(v5p/v6e/v7),VMREG 从 8 翻倍到 16,而 PREG 实际上从 15 缩小到 14。为 v5+ target 编号寄存器的重新实现者如果假定 v2 时代的文件大小,就会放错每个 Sreg、Vmreg 和 Vreg 的
arch_regno,因为块基址会随按类别计数移动。
读取路径
ToArchRegno —— 反向查找
function ToArchRegno(this, arch_regno): // sub_1275f580
if arch_regno == 0: Fatal("reg_num != kNoRegister")
if arch_regno >= this.total_registers_: Fatal("reg_num < total_registers_")
if arch_regno >= this.idx_to_regno_.size: BUG()
return this.idx_to_regno_[arch_regno] // [this+0x80][arch_regno] -> (RegisterType, regno)ToRegNum —— 正向查找
function ToRegNum(this, type, regno): // sub_1d5a9000
if type == kNone: Fatal("register_type != RegisterType::kNone")
key = (regno << 32) | type
it = this.regno_to_idx_.find(key) // [this+0x98]
if not found: return 0
assert it.value <= 255 // "iter->second <= UINT8_MAX"
return it.value // arch_regnoToArchRegString —— 可打印名称
ToArchRegString @ 0x1275e2a0 会把一个 slot 解析为其可打印形式:在 (type, regno) = ToArchRegno(arch_regno) 之后执行 StrCat(RegisterTypeToMnemonic(type), FastIntToBuffer(regno))。因此架构寄存器 arch_regno 会打印为 "<p|s|vm|v><regno>",例如 v3、s12、vm5、p2。
| 函数 | 地址 | 读取 |
|---|---|---|
ToArchRegno | 0x1275f580 | [this+0x80] 反向向量 |
ToRegNum | 0x1d5a9000 | [this+0x98] 正向 map |
ToArchRegString | 0x1275e2a0 | ToArchRegno + 助记符 |
陷阱 ——
ToRegNum在未命中时返回 0(kNoRegister哨兵),而不是 fatal,并且会断言找到的arch_regno适合放入uint8。某个(type, regno)从未为此 target 编号的虚拟寄存器会静默解析为“无寄存器”。相反,ToArchRegno会在arch_regno == 0时 fatal。两个方向的错误处理并不对称。
按 Opcode 的 Variant 分类
ComputeXluOperations 选择的每个 XLU opcode(见 ResultFifo 和 ArchRegister 枚举)都会被包进 std::variant<TransposeTile, RpuOperation, XluControlOperation>。该 variant 宽 0x48 字节(步长由 lea (idx,idx,8); shl 3 = ×72 确认),判别字节位于 +0x40(emitter 写入 movb $0/$1/$2, 0x40(...))。索引由四个字节精确的分类谓词选择:
| 分类器 | 地址 | 谓词 | 选择 | Variant idx |
|---|---|---|---|---|
LloOpcodeUsesTranspose | 0x1d60bda0 | (op-0xa6)<2 || (op-0x154)<2 | {0xa6, 0xa7, 0x154, 0x155} | 0 TransposeTile |
LloOpcodeUsesRpu | 0x1d60c2c0 | (op-0xf5)<0xd (加上下方的低段 bitmask) | {0xf5..0x101} | 1 RpuOperation |
LloOpcodeIsRpuControl | 0x1d60c1e0 | (op-0x8b)<2 | {0x8b, 0x8c} | 2 XluControlOperation |
LloOpcodeIsRpuResult | 0x1d60c420 | (op-0x14f)<2 | {0x14f, 0x150} | 2 XluControlOperation |
Variant 主体
TransposeTile(索引 0)是唯一拥有资源的备选项。布局:+0x00处是InlinedVector<ArchRegister,2>读集合,+0x18处是InlinedVector<ArchRegister,2>写集合,+0x30处是一个i64partner/tile 字段,+0x37处是一个打包i32,判别符在+0x40。当两个堆后备寄存器集合的大小超过 2 的内联容量时,析构函数会释放它们。RpuOperation(索引 1)可平凡复制(≤0x28 字节)。它在+0x10持有一个LloValue*,并带有内联 fusion-key 元数据(RpuOperationMetadata:+0x00处u16 opcode,+0x08处i64pattern/segment key,+0x10处optional<i64>,由[+0x18]==1门控)。元数据相等的两个RpuOperation是 fusion 候选。XluControlOperation(索引 2)可平凡析构:+0x00处一个LloValue*加内联 control 元数据,没有拥有资源的 vector。
LloOpcodeUsesRpu 还有一个次级低 opcode 分支:对于 op <= 0x3b,它测试 _bittest64(0xC40000000000000, op),该值设置了 {0x36, 0x3a, 0x3b} 这些 bit(位于跨 lane reduce 段之外的 RPU op)。主分支 (op-0xf5)<0xd 覆盖 13 个 reduce op。
CrossXlu 数据依赖跟踪器
ComputeXluOperations 附加到每个 variant 的按 op ArchRegister 读/写集合并不由编号本身使用,而是驱动一个跨 lane op 的列表调度器。CrossXluOperationsDataDependencyTracker::Create @ 0x126cc9a0 按拓扑顺序遍历 LloDependencyGraph,标记其 XLU-operation variant 在输入集合中的节点,分配一个 0x330 字节的 tracker(operator new(0x330)),并把每个 XLU op 连接到与它有真实数据依赖的其他 XLU op。
function CrossXlu::Create(graph, xlu_ops, reverse): // sub_126cc9a0
nodes = graph.NodesInTopologicalOrder(reverse)
// pass 1: mark nodes whose variant is in xlu_ops (flat_hash_set::find)
for node in nodes:
if xlu_ops.contains(node.value.variant): node.is_xlu = true
tracker = new CrossXluOperationsDataDependencyTracker(region) // 0x330 bytes
// pass 2: per marked node, build cross-XLU edges
for node in marked_nodes:
tracker.internal_graph.AddUnsequencedNode(node.instruction) // [tracker+0x100]
for each def-use edge:
if LloDependencyGraphNode::IsOperandOf(other): // sub_14427d60
add edge (producer -> consumer); ++consumer.in_edges
return StatusOr<unique_ptr<...>>tracker 构造函数(0x126cda80)在 tracker+0x100 构建一个内部 LloDependencyGraph,并在 tracker+0xf8 构建一个 LatencyTable(从 TpuVersion 和 Target+0x398 字段创建)。就绪性是标准的入边计数测试:
function XluOperationIsReady(tracker, op): // sub_126cd920
// dispatch on the variant discriminant at op[0x40]:
switch op[0x40]:
case 0: value = TransposeTile.read_set.last_element // InlinedVector at op+0x08
case 1: value = RpuOperation.lloValue // op+0x10
case 2: value = XluControlOperation.lloValue // op+0x00
default: Fatal("control != nullptr")
node = tracker.internal_graph.GetNode(value) // [tracker+0x100]
return node.in_edge_count == 0 // all predecessor XLU ops scheduledRemoveScheduledXluOperation @ 0x126ccfa0 是“触发”步骤:它移除一个已调度 op,并递减其后继的入边计数。这些边正是 LloValue def-use 链:某个生产者 XLU op 写入的 ArchRegister 被后续 XLU op 读取时,它会成为后者的前驱(IsOperandOf 确认真实的 RAW/WAR 关系)。tracker 的消费者包括 ComputeCombinablePairs(0x126d2480,融合拥有相同 RpuOperationMetadata 的 op)、ReorderToShortenCriticalPath(0x126d3460)、AssignXlu(0x126d3100)和 AssignSourceBus(0x126d70e0)。
| 函数 | 地址 | 角色 |
|---|---|---|
CrossXlu...::Create | 0x126cc9a0 | 拓扑标记 → new tracker → 按节点建边 |
CrossXlu...::ctor | 0x126cda80 | 内部 LloDependencyGraph + LatencyTable |
XluOperationIsReady | 0x126cd920 | 就绪性 = 入边计数 == 0 |
RemoveScheduledXluOperation | 0x126ccfa0 | 列表调度器触发步骤 |
LloDependencyGraphNode::IsOperandOf | 0x14427d60 | 确认 RAW/WAR 数据依赖 |
注意 —— 依赖跟踪器馈入关键路径重排的边权重是
LatencyTable中按版本的 op 延迟(LatencyTable::Create@0x1c89fba0)。
相关组件
| 组件 | 关系 |
|---|---|
ArchRegister / ArchRegisterInstance(result-fifo / arch-register 枚举) | 此编号映射到 (RegisterType, regno) 的物理 slot 层 |
ComputeXluOperations(result-fifo / arch-register 枚举) | 构建 tracker 使用的按 op ArchRegister 读/写集合 |
Target 寄存器计数字段 | 编号做前缀和的四个按类别计数(0x498..0x4c0) |
交叉引用
- ResultFifo 和 ArchRegister 枚举 —— 50 成员
ArchRegister物理编号,以及此编号所在层之下的 25 个 result FIFO - XLU Op 名册 —— 其 op 会被 variant 分类器标记的跨 lane opcode→factory 表
- Slot: VPU —— 此编号所分配寄存器所属的 vector-processing slot
- Slot: MXU —— matmul slot,其 gain/latch matrix-register file 不同于这些 GPR 类别
- MC Emitter —— 在编码时消耗
arch_regno的机器码 emitter - opcode_info_big 记录格式 —— 持有
ArchRegister读/写列表的按 opcode 描述符 - Bundle 模型 —— 这些架构寄存器编码进入的按代际 VLIW bundle