Skip to content

ArchRegno 编号

本页中的每个偏移、值和地址都按字节精确读取自 libtpu-0.0.40-cp314 wheel 中的 libtpu.so(BuildID md5 89edbbe81c5b328a958fe628a9f2207d)。其他版本会有所不同。

摘要

TensorCore 编码器在发出一个 bundle slot 时,需要为每个架构寄存器取得一个单一的稠密整数,但 LLO IR 携带的是由 RegisterTypepreg / sreg / vmreg / vreg)加该类别内 regno 描述的虚拟寄存器。桥接层是 RegisterNumbering:一个按 Target 创建的对象,为每个 (RegisterType, regno) 分配一个连续的架构寄存器编号arch_regno),同时维护正向映射((type, regno) → arch_regno,由 ToRegNum 读取)和反向向量(arch_regno → (type, regno),由 ToArchRegno 读取)。由于不同硅片代际的寄存器文件大小不同,编号会在每次构造 Target 时根据四个代际特定的计数字段重新构建。本页记录这个构建过程,即 Target::InitRegisterNumberingRegisterNumbering::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 thisSequencer 角色存在的类别
1Target + 0x008Primary / TensorCore全部四类(preg, sreg, vmreg, vreg)
2Target + 0x138非 primary(仅 preg+sreg)preg, sreg
3Target + 0x268非 primary(仅 vmreg+vreg)vmreg, vreg

步长 0x130 就是 RegisterNumbering 大小(0x138 - 0x008 = 0x268 - 0x138 = 0x130)。反编译确认了在 this+8this+3120x138)和 this+6160x268)处的三次 RegisterNumbering::Init 调用。

算法

每个映射都是一个新构建的 absl::flat_hash_map<RegisterType, pair<int count, string name>>,含四个条目(键 1..4),传给 RegisterNumbering::Init。按类别的 count 来自 Target 寄存器计数字段块;按类别的 name 来自 TargetEnvironment

c
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::InitRegisterNumbering0x1d614200构建 3 个 RegisterNumbering 对象
Target::RegisterCount0x1d617120同一组 4 个计数字段的读取侧,按 (seq, type) 取键
Target::SregCount0x1d6152c0mov 0x498(rdi),eax —— 确认 Sreg 偏移
Target::VregCount0x1d6152e0mov 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,同时构建映射的两个方向。

对象布局

c
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 开始,步长为 0x18AddRegister 计算 &this[8 + 0x18*type]);五个类型占用 0x08..0x80。反向向量位于 +0x80int* 索引 32),正向 map 位于 +0x98int* 索引 38),二者都由 AddRegisterpush_backinsert 目标逐字节确认。

算法

c
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

c
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::Init0x1d622520前缀和 + 双遍分配
RegisterNumbering::AddRegister0x1d622bc0分配一个 arch_regno,构建两种映射
RegisterTypeToMnemonic0x1d640600名称的类别前缀(p/s/vm/v
FastIntToBuffer0x211719e0为 regno 后缀做整数→ASCII
RangeSpec::Match0x1d624c80配置驱动的按 regno 过滤器(设置 is_pseudo
~RegisterNumbering0x1d491e60释放 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 的寄存器名称/范围过滤器;该机制返回传给 AddRegisteris_pseudo 标志。因此名义类别计数可能大于可用(可分配)计数。


按代际的寄存器文件大小

编号所消耗的四个计数字段由每个 TpuVersion 的 chip-parts sequencer 描述符填充。字段偏移和前缀和布局是字节精确的;下面的数值计数是按代际取键的 chip-parts 值。

代际(代号)TpuVersionSREG (0x498)VREG (0x49c)VMREG (0x4a0)PREG (0x4a4)Arch regs(近似)
v2 / v3 / v4 (jellyfish / dragonfish / pufferfish)kJellyfish=0, kDragonfish=1, kPufferfish=23232815~88 (+1 sentinel)
v5p (+v5e lite) / v6e / v7 (viperfish / ghostlite / 6acc60406)kViperfish=3, kGhostlite=4, gen-5=532641614~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/0x4a4SregCount/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 —— 反向查找

c
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 —— 正向查找

c
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_regno

ToArchRegString —— 可打印名称

ToArchRegString @ 0x1275e2a0 会把一个 slot 解析为其可打印形式:在 (type, regno) = ToArchRegno(arch_regno) 之后执行 StrCat(RegisterTypeToMnemonic(type), FastIntToBuffer(regno))。因此架构寄存器 arch_regno 会打印为 "<p|s|vm|v><regno>",例如 v3s12vm5p2

函数地址读取
ToArchRegno0x1275f580[this+0x80] 反向向量
ToRegNum0x1d5a9000[this+0x98] 正向 map
ToArchRegString0x1275e2a0ToArchRegno + 助记符

陷阱 —— 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
LloOpcodeUsesTranspose0x1d60bda0(op-0xa6)<2 || (op-0x154)<2{0xa6, 0xa7, 0x154, 0x155}0 TransposeTile
LloOpcodeUsesRpu0x1d60c2c0(op-0xf5)<0xd (加上下方的低段 bitmask){0xf5..0x101}1 RpuOperation
LloOpcodeIsRpuControl0x1d60c1e0(op-0x8b)<2{0x8b, 0x8c}2 XluControlOperation
LloOpcodeIsRpuResult0x1d60c420(op-0x14f)<2{0x14f, 0x150}2 XluControlOperation

Variant 主体

  • TransposeTile(索引 0)是唯一拥有资源的备选项。布局:+0x00 处是 InlinedVector<ArchRegister,2> 读集合,+0x18 处是 InlinedVector<ArchRegister,2> 写集合,+0x30 处是一个 i64 partner/tile 字段,+0x37 处是一个打包 i32,判别符在 +0x40。当两个堆后备寄存器集合的大小超过 2 的内联容量时,析构函数会释放它们。
  • RpuOperation(索引 1)可平凡复制(≤0x28 字节)。它在 +0x10 持有一个 LloValue*,并带有内联 fusion-key 元数据(RpuOperationMetadata+0x00u16 opcode+0x08i64 pattern/segment key,+0x10optional<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。

c
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(从 TpuVersionTarget+0x398 字段创建)。就绪性是标准的入边计数测试:

c
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 scheduled

RemoveScheduledXluOperation @ 0x126ccfa0 是“触发”步骤:它移除一个已调度 op,并递减其后继的入边计数。这些边正是 LloValue def-use 链:某个生产者 XLU op 写入的 ArchRegister 被后续 XLU op 读取时,它会成为后者的前驱(IsOperandOf 确认真实的 RAW/WAR 关系)。tracker 的消费者包括 ComputeCombinablePairs0x126d2480,融合拥有相同 RpuOperationMetadata 的 op)、ReorderToShortenCriticalPath0x126d3460)、AssignXlu0x126d3100)和 AssignSourceBus0x126d70e0)。

函数地址角色
CrossXlu...::Create0x126cc9a0拓扑标记 → new tracker → 按节点建边
CrossXlu...::ctor0x126cda80内部 LloDependencyGraph + LatencyTable
XluOperationIsReady0x126cd920就绪性 = 入边计数 == 0
RemoveScheduledXluOperation0x126ccfa0列表调度器触发步骤
LloDependencyGraphNode::IsOperandOf0x14427d60确认 RAW/WAR 数据依赖

注意 —— 依赖跟踪器馈入关键路径重排的边权重LatencyTable 中按版本的 op 延迟(LatencyTable::Create @ 0x1c89fba0)。


相关组件

组件关系
ArchRegister / ArchRegisterInstanceresult-fifo / arch-register 枚举此编号映射到 (RegisterType, regno) 的物理 slot 层
ComputeXluOperationsresult-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