Skip to content

注册表中介标志

本页所有地址均适用于 libtpu-0.0.40-cp314 wheel 中的 libtpu.so(构建 libtpu_lts_20260413_b_RC00,build-id md5 89edbbe81c5b328a958fe628a9f2207d,781,691,048 字节,ELF x86-64 DYN,未 strip;反修饰后的 C++ 符号按原文引用)。其他版本会有所不同。

摘要

大多数 TPU 编译旋钮的读取方式如 autoproto-autoor-resolution.md 所述:生成的访问器从 TpuCompilationEnvironment(TCE)结构体的固定偏移加载一个 AutoProto*,调用 AutoOr<T>::FromProtoOrDie,并应用极性测试。该路径是偏移键控的,也就是字段在结构体中的位置会在编译时固化进访问器。还有一小类标志不是这样读取的。它们的消费者绝不会从固定偏移加载字段,而是在运行时按名称查找值,通过由 flag↔field 注册表中介的 TCE proto 反射读取来完成。本页负责说明这条注册表中介的读取路径:按名称反射读取器 GetFieldValueIfNotDefault<T> @ 0x1c6f1a80FlagFieldMappings flag↔field 桥 @ 0x2257ef50、按 TpuVersion 划分的 LegacyEvictions/LegacyPrefetches 同步标志注册表、写入侧 SetFieldFromFlagString @ 0x1d73fcc0,以及一个标志驱动编译器变换的完整示例:xla_tpu_impure_use_iteration_maskRaggedDotExpander iteration-mask 降低。

参照框架是 XLA 的 xla::DebugOptions proto 后端标志,但带有一个 libtpu 特有的细节,重新实现者必须复现。CLI 标志(absl::CommandLineFlag,即 .data 中的 FLAGS_<name> 对象)和 TCE proto 字段是运行时桥接的两个独立事物FlagFieldMappingsNoDestructor 构造函数遍历每个 TCE proto 字段,调用 absl::FindCommandLineFlag(field_name),并将配对插入两个 FlatHashMap(flag→field 和 field→flag)。在 env 组装时,SetFieldFromFlagString 将标志解析回其 FieldDescriptor,并把解析出的值写入该 proto 字段;在读取时,GetFieldValueIfNotDefault<T> 按名称查找字段,通过反射读取它,并且仅在它不同于全默认 env 时返回。标志值的流向是 CLI 字符串 → TCE proto 字段 → 按名称反射读取;消费者处从不会出现 lea FLAGS_<name>。正是这层间接性让一个标志成为“注册表中介”的:消费者绑定到名称,而不是结构体偏移或 FLAGS_ 符号。

这种设计带来的第三个尖锐后果是,标志可以已注册但未接线enable_lem_schedulerexplicit_evict_memory_limit_kib 都已注册(因此 --helpfull 会列出它们,CLI 也能解析它们),但没有静态可达的消费者:它们的 FLAGS_ 对象恰好只有一个代码 xref,即注册用的 lea;并且它们的名称不存在于序列化的 TCE FileDescriptorProto 中,所以按名称读取器的 FindFieldByName 会失败。它们是残留物。本页结构如下:三种读取类别;按名称反射读取器和桥;证明这两个具名标志未接线,并说明实际驱动这些行为的机制;最后是完整的 RaggedDot flag→transform 示例。

对于重新实现,契约是:

  • 三种读取类别 — (A) FLAGS_ 锚定的直接读取(lea FLAGS_FlagImpl+0x58AutoOr 内联解析),(B) 按名称 TCE 反射(注册表中介路径),(C) 已注册但未接线(名称不是 proto 字段)。重新实现者必须通过消费者的读取指令区分它们,而不是通过标志名称或类型。
  • 按名称读取机制GetFieldValueIfNotDefault<T> = FindFieldByNameGetFieldValue 反射 → 与全默认 env 求差 → optional<T>FlagFieldMappings 双向 FlatHashMap 桥;写入侧 SetFieldFromFlagString
  • 完整 flag→transform 示例impure_use_iteration_mask(读取类别 A,AutoOr<bool>,AUTO=ON,TpuVersion>=3 门控)→ RaggedDotExpanderIota + Compare + AND + Select(mask, conv, 0) iteration mask,用于在 ragged convolution 中将跨组乘积清零。
按名称读取器GetFieldValueIfNotDefault<long> @ 0x1c6f1a80FindFieldByName + TpuCompEnvReflection::GetFieldValue + 与默认值求差
Flag↔field 桥FlagFieldMappings::GetInstance @ 0x2257ef50 · 构造函数 @ 0x1d753ce0(双向 FlatHashMap,每字段一次 FindCommandLineFlag
写入侧SetFieldFromFlagString @ 0x1d73fcc0GetFieldForFlag + ParseFlagFromString 写入 proto 字段
旧式同步标志注册表LegacyEvictionsFlagRegistry @ 0x22579860 · LegacyPrefetchesFlagRegistry @ 0x225798a0(按 TpuVersion 划分的 StaticMap
MSA 消费者ComputeMemoryManagementSflagUsage @ 0x1c6f1580 — 2× GetFieldValueIfNotDefault<long>@ 0x1c6f185e/0x1c6f1888
未接线标志FLAGS_xla_tpu_enable_lem_scheduler @ 0x223c72a8 · FLAGS_xla_tpu_explicit_evict_memory_limit_kib @ 0x223c50d0 — 每个 1 个 xref(仅注册)
注册器_GLOBAL__sub_I_tpu_compilation_environment.cc @ 0x2135cba0..0x21360ef0(每个标志一次 RegisterCommandLineFlag @ 0x21114cc0
示例标志FLAGS_xla_tpu_impure_use_iteration_mask @ 0x223a7dd0AutoOr<bool>,FlagOps 0x1d6b5840,6 个 xref(5 个消费者 + 注册器)
示例变换RaggedDotExpander::RunImpl @ 0x10fae060CreateOutputMask @ 0x10fb2900CreateConvolutionSelectFusion @ 0x10fb31e0
置信度CONFIRMED(字节锚定 vs 反编译),除非某行或提示另有说明

1. 三种读取类别

目的

TCE 编译旋钮的可以通过三条不同路径到达消费者。这一区分很重要,因为它会同时改变重新实现者查找消费者的位置,以及该标志是否在运行时真的有影响。对每个标志而言,这些类别互斥,并且标志所属类别只能从其消费者的读取指令中看出来。

类别

类别消费者处的读取指令示例标志效果
(A) FLAGS_ 锚定的直接读取lea FLAGS_<name>(%rip)FlagImpl::ReadOneWord/ReadOneBoolFlagImpl+0x58 缓存)→ AutoOr 极性内联impure_use_iteration_mask, impure_enable_masked_fusion_iteration_skipper存活
(B) 按名称 TCE 反射没有 lea FLAGS_GetFieldValueIfNotDefault<T>(name, env)FindFieldByName + GetFieldValue 反射xla_jf/vf/zf/gf_vmem_max_outstanding_evictions/_prefetches, explicit_prefetch_memory_limit_kib存活
(C) 已注册但未接线无 — FLAGS_ 对象唯一的 xref 是注册 lea;名称不是 TCE proto 字段enable_lem_scheduler, explicit_evict_memory_limit_kib在 v0.0.40 中死亡

类别 (A) 是 autoproto-autoor-resolution.mdflag-families.md 为大部分旋钮表面记录的路径:一段生成的访问器直接从 FLAGS_ 对象或 TCE 结构体偏移加载值。类别 (B) 是本页主题。类别 (C) 是退化的 (B):读取机制存在,但标志名称从未加入 TCE proto,因此按名称读取器永远找不到它。

陷阱 — 类别 (B) 和 (C) 在 FLAGS_ 层面共享同一个静态签名:恰好一个代码 xref,即注册 lea。一次覆盖整个 .text 的 RIP-relative lea/mov 扫描无法区分二者,因为二者都没有引用 FLAGS_ 对象的消费者。判别器是序列化的 TCE proto:类别 (B) 标志的名称是 proto 字段 token(所以 FindFieldByName 成功);类别 (C) 标志的名称不是。只按 xref 数量分类的重新实现者会把每个未接线标志都误归为单纯的注册表中介标志。

标志值如何进入 proto(写入侧)

对于类别 (B),标志值必须先写入 TCE proto 字段,随后任何按名称读取才能观察到它。这个桥是 SetFieldFromFlagString @ 0x1d73fcc0

c
function SetFieldFromFlagString(CommandLineFlag* flag, string_view value,    // 0x1d73fcc0
                                TpuCompilationEnvironment* env):
    field = TpuCompEnvReflection::GetFieldForFlag(flag)    // flag → FieldDescriptor(field->flag 映射)
    if !field.ok():                                        // FATAL "Flag is not found for field" 路径
        return Status(... "tpu_compilation_environment.cc":5905)
    TpuCompEnvReflection::ParseFlagFromString(flag, field, value)  // 解析 + 通过反射写入 env[field]
    ...
```text

写入会通过桥(§2)把标志解析回其 `FieldDescriptor`,然后 `ParseFlagFromString` 根据字段类型解析 CLI 字符串,并通过反射将其写入 proto。之后,`GetFieldValueIfNotDefault<T>(field_name, env)` 就会看到非默认值。

> **说明 —** 写入侧已做结构确认:`SetFieldFromFlagString` 调用 `GetFieldForFlag`,然后调用 `ParseFlagFromString`,形状正是 flag→field 后写入(CONFIRMED)。此处未逐一重新走查每个有类型字段的逐标志解析语法(它是 [autoor-parse-grammar.md](autoor-parse-grammar.md) 上的 `AutoOr`/`Tristate` 解析语法)。

---

## 2. 按名称反射读取器和桥

### 目的

类别 (B) 读取是一种*反射*读取:消费者持有字段*名称*(一个 `string_view`),而不是结构体偏移,并在运行时针对 TCE proto 的描述符解析该值。它由两部分实现:读取器 `GetFieldValueIfNotDefault<T>`,以及在标志和字段之间双向映射的 `FlagFieldMappings` 桥。

### 入口点

```text
<consumer>(name, env)                                    ── 例如 ComputeMemoryManagementSflagUsage 0x1c6f1580
  └─ GetFieldValueIfNotDefault<T>(name, env)             ── 0x1c6f1a80(T=long 实例化)
       ├─ TpuCompilationEnvironment::GetMetadata()->descriptor   ── 0x1c6f1aa0
       ├─ Descriptor::FindFieldByName(name)              ── 0x20e57900;null → "... not found ..." 错误
       ├─ TpuCompEnvReflection::GetFieldValue(field, env)        ── 0x1d7523a0 → 覆盖每个 AutoOr arm 的 variant
       ├─ GetFieldValue(field, GetTpuCompEnvWithDefaultValues())  ── 默认 env,0x1d73f100
       └─ __fmatrix variant visitor: (val != def) ? optional<T>(val) : nullopt

算法

long 实例化 @ 0x1c6f1a80 是规范实现体;其他有类型实例化只在所提取的 variant arm 上不同:

c
function GetFieldValueIfNotDefault<long>(string_view name,                   // 0x1c6f1a80
                                         ObjectView<TCE> env):
    desc = TpuCompilationEnvironment::GetMetadata()->descriptor   // 0x1c6f1aa0(来自 TCE_globals_)
    field = desc->FindFieldByName(name)                           // 0x1c6f1ab5
    if !field:                                                    // 不是 proto 字段
        return StatusError(StrCat(name, " ... not found ..."))    // 0x1c6f1c27 / StrCat @ 0x21174860
    val = TpuCompEnvReflection::GetFieldValue(field, env)         // 0x1c6f1ad3 → absl::variant
    def = TpuCompEnvReflection::GetFieldValue(field,
              GetTpuCompEnvWithDefaultValues())                   // 0x1c6f1b36 / 0x1c6f1b48
    // __fmatrix visitor 逐元素比较两个 variant                         // 0x1c6f1bdb
    if val == def:
        return nullopt                                            // "IfNotDefault" 语义
    return optional<long>(extract_long(val))
```text

`GetFieldValue` 返回的 variant 是对 TCE 旋钮可携带的**所有**值类型的求和类型:反编译中 visitor 的类型列表是 `a h i l j m d f b`(bool/int8…/int64/uint64/double/float),外加 `RangeSpecProto`、`RepeatedStrings`、`SparseDenseMatmulFdoConfig`、`SlicedPrefetchOptions`、`MemoryBoundLoopOptimizerOptions`、`PreferredPrefetchOverrides`、`MsaSortOrderOverrides`、`BufferContentsSanitizerConfig`、`BufferIsolationConfig` 和 `AutoProto`。因此三态旋钮的 `AutoProto` oneof 在这条路径上是*通过反射*解析的,而不是通过内联的 `AutoOr<T>::FromProtoOrDie` 带解析的;存储单元相同(§ [autoproto-autoor-resolution.md](autoproto-autoor-resolution.md)),读取器不同。

> **怪癖 —** 该读取是差分读取,不是绝对读取。当字段等于 `GetTpuCompEnvWithDefaultValues()` 中的值时,`GetFieldValueIfNotDefault` 返回 `nullopt`,而不是某个哨兵值或已存储值。想知道“用户是否设置了这个旋钮?”的消费者会得到真实答案;想知道“包含默认值在内的有效值是什么?”的消费者必须在 optional 为空时提供自己的默认值。MSA 消费者(§3)正是这样使用它的:除非字段不同于默认值,否则忽略该字段。

### flag↔field 桥

标志和 proto 字段在运行时由 `FlagFieldMappings` 连接,它是一个只构造一次的 `NoDestructor` 单例:

```c
function FlagFieldMappings::ctor():                              // 0x1d753ce0(NoDestructor 主体)
    desc = TpuCompilationEnvironment::GetMetadata()->descriptor  // 0x1d753...(GetMetadata @ 0x1db635c0)
    n    = desc->field_count                                     // *(int*)(meta + 8)
    for i in 0 .. n-1:
        field = desc->field(i)                                   // 字段步幅遍历,meta+64 基址
        flag  = absl::FindCommandLineFlag(field->name())         // 0x1d753dbf(FindCommandLineFlag @ 0x21115120)
        if flag:
            flag_to_field[flag]  = field                         // 0x1d753e1f  FlatHashMap<CommandLineFlag*, FieldDescriptor*>
            field_to_flag[field] = flag                          // 0x1d753e7f  FlatHashMap<FieldDescriptor*, CommandLineFlag*>

配套访问器 TpuCompEnvReflection::GetFlagForField @ 0x1d74ad40 读取 field→flag 映射,并在某字段没有已注册标志时 FATAL("Flag is not found for field" 字符串 @ 0x86f971b)。

陷阱 — 该桥是通过迭代proto 字段并为每个字段查找标志来构建的,而不是迭代标志。因此该桥只覆盖名称是 TCE proto 字段的标志。名称不是字段的已注册标志(类别 C)永远不会插入任一映射;若被请求,GetFlagForField 会 FATAL,而按名称读取器的 FindFieldByName 会返回 null。桥的定义域就是类别 (B) 标志集合。

函数映射

函数地址作用
GetFieldValueIfNotDefault<long>0x1c6f1a80按名称反射读取器 — FindFieldByName + 与默认值求差
Descriptor::FindFieldByName0x20e57900名称 → FieldDescriptor(null ⇒ 不是字段)
TpuCompEnvReflection::GetFieldValue0x1d7523a0反射读取 → 覆盖所有 arm 类型的 variant
GetTpuCompEnvWithDefaultValues0x1d73f100读取差分所针对的全默认 env
FlagFieldMappings ctor0x1d753ce0构建双向 FlatHashMap(flag↔field)
FlagFieldMappings::GetInstance()::mappings0x2257ef50NoDestructor 单例存储(访问器被内联)
TpuCompEnvReflection::GetFlagForField0x1d74ad40field→flag 查找(缺失时 FATAL)
SetFieldFromFlagString0x1d73fcc0写入侧 — 标志字符串 → 通过反射写入 proto 字段
RegisterCommandLineFlag0x21114cc0注册器中每个标志的注册调用

3. 旧式同步标志注册表(类别 (B) 示例)

目的

最干净的存活类别 (B) 消费者是 MSA(memory-space assignment)逐出/预取限制读取器。逐出和预取的 outstanding-op 限制是按 TPU 系列划分的 TCE 字段,带有系列特定名称(Jellyfish 用 xla_jf_…,还有 xla_vf_…xla_zf_…xla_gf_…)。消费者不是为每个系列硬编码一个偏移,而是在按 TpuVersion 划分的注册表中查找字段名称,然后按名称读取该字段。

算法

c
function ComputeMemoryManagementSflagUsage(ObjectView<TCE> env, Target& target):  // 0x1c6f1580
    ver = target.tpu_version
    evict_name    = LegacyEvictionsFlagRegistry[ver]    // StaticMap<TpuVersion, string_view> @ 0x22579860
    prefetch_name = LegacyPrefetchesFlagRegistry[ver]   //                                    @ 0x225798a0
    evict_opt     = GetFieldValueIfNotDefault<long>(evict_name, env)       // 0x1c6f185e
    prefetch_opt  = GetFieldValueIfNotDefault<long>(prefetch_name, env)    // 0x1c6f1888
    // "legacy_non_default_value"(str @ 0x94adace):仅当字段不同于默认值时才使用该字段
    ...将 evict_opt / prefetch_opt 应用于 sync-flag usage...
```text

这两个注册表都是以 `tpu::TpuVersion` 为键的 `util_registration::StaticMapBase` 单例(反编译中已验证:`StaticMapBase<…LegacyEvictionsFlagRegistry, tpu::TpuVersion, string_view, …>::GetSingleton`)。它们在 `_GLOBAL__sub_I_sync_flag_util_{dragonfish,jellyfish,pufferfish,viperfish,ghostlite}.cc` 中按系列通过 `InsertValue @ 0x1c6f8760`/`0x1c6f8940` 填充。

### 已注册字段名称

| 注册表 | 字段名称模式(按系列) | 锚点 |
|---|---|---|
| `LegacyEvictionsFlagRegistry @ 0x22579860` | `xla_{zf,vf,jf,gf}_vmem_max_outstanding_evictions`(+ `cmem` 变体) | `"xla_jf_vmem_max_outstanding_evictions" @ 0x853d9d7` |
| `LegacyPrefetchesFlagRegistry @ 0x225798a0` | `xla_{zf,vf,jf,gf}_vmem_max_outstanding_prefetches` | `"xla_jf_vmem_max_outstanding_prefetches" @ 0x8569c57` |

所有这些名称都是 TCE proto 字段(存在于序列化的 `descriptor_table_protodef_AiDtBo5TtCO` 中),所以 `FindFieldByName` 成功,并且当用户设置标志时,差分读取会返回真实值。这就是规范的注册表中介路径:注册表将运行时键(`TpuVersion`)映射到字段*名称*,值再按该名称读取。

---

## 4. 未接线的一对标志(类别 (C))

### 目的

`enable_lem_scheduler` 和 `explicit_evict_memory_limit_kib` 是已记录的类别 (C) 例外:已注册,CLI 可解析,`--helpfull` 会列出,但在 v0.0.40 中没有静态可达的值消费者。证明来自多个角度且精确到字节;本节记录它,避免重新实现者浪费时间接线一个死亡标志。

### 证明

```text
1. 覆盖整个 .text 的 RIP-relative lea/mov 扫描(0xe635560..0x213f81a0):
     FLAGS_xla_tpu_enable_lem_scheduler              @0x223c72a81 xref: lea @0x21360e44
     FLAGS_xla_tpu_explicit_evict_memory_limit_kib   @0x223c50d01 xref: lea @0x21360976
   每个 xref 都是 `lea FLAGS_…(%rip),%rdi ; mov %rbx,%rsi ; call RegisterCommandLineFlag@0x21114cc0`
   位于注册器 _GLOBAL__sub_I_tpu_compilation_environment.cc(0x2135cba0..0x21360ef0)内部。
   (对比:impure_use_iteration_mask @0x223a7dd06 个 xref = 5 个消费者 + 1 个注册器。)
2. 名称字符串 xref 扫描:0 个 `lea name-string` 站点 → 任何调用点都没有按名称 FindCommandLineFlag。
3. 序列化的 TCE FileDescriptorProto descriptor_table_protodef_AiDtBo5TtCO @0xbf9d5d0 包含
   "TpuCompilationEnvironment""mxu_latency""AutoProto"
   "explicit_prefetch_memory_limit_kib"(@protodef+0x1a45f)— 但不包含 "enable_lem_scheduler"
   也不包含 "explicit_evict_memory_limit_kib"。⇒ FindFieldByName 失败 ⇒ 按名称读取器无法触达它们。

每个标志的 FLAGS_ 对象都像其他所有 libtpu 标志一样由 .rela.dyn 填充重定位:+0x08 名称字符串,+0x10 类型名字符串,+0x20 FlagOps(AutoOr TypeId),+0x28 HelpGen;其默认生成器发出 AUTO(enable_lem_schedulermovw $0 ⇒ AUTO;explicit_evictmovq $0; movb $0,0x8{0, has0} ⇒ AUTO)。它们是形态良好但没有读取器的标志。

标志FLAGS_类型唯一 xref实际驱动行为的机制
xla_tpu_enable_lem_scheduler0x223c72a8AutoOr<bool>(FlagOps 0x1d6b5840lea @ 0x21360e44(注册)运行时谓词 ModuleContainsLEMSparseCoreInstruction @ 0x13853280
xla_tpu_explicit_evict_memory_limit_kib0x223c50d0AutoOr<int64_t>(FlagOps 0x1d700120lea @ 0x21360976(注册)旧式逐版本 vmem-evict TCE 字段(§3)

实际驱动 LEM scheduler 的机制

LEM-scheduler 决策不是该标志。在 RunHloScheduler @ 0x1096fac0 中:

c
function RunHloScheduler(module, env, target):                  // 0x1096fac0
    use_lem = offloader_util::ModuleContainsLEMSparseCoreInstruction(module)   // 0x13853280 → r8b
    cfg     = GetSchedulerConfig(module, env, target, /*bool=*/use_lem)        // 0x10974aa0
    est     = GetLatencyEstimator(target, cfg, env)             // 0x10974e00
              // → CostModelLatencyEstimator @ 0x10ff8a60  ("LEM" cost estimator)
```text

`ModuleContainsLEMSparseCoreInstruction` 扫描 `SparseCoreDataFormatOffloader` / `SparseCoreOffloadingOptions_OffloadFeature` 指令(`@ 0x21925478`/`0x21925488`)。这里的 "LEM" 是 Large-Embedding-Model SparseCore offload,而 LEM-aware scheduler 是由模块*包含*此类指令来选择的,这才是该残留标志原本想扮演的角色。

> **说明 —** 逐出限制的不对称性是最尖锐的信号:`explicit_prefetch_memory_limit_kib` 是 TCE proto 字段(`@ protodef+0x1a45f`,类别 B),但它的孪生 `explicit_evict_memory_limit_kib` 不是字段(类别 C)。二进制中无法恢复 evict 字段究竟是被移除(弃用)还是从未合入;两者都会表现为“未接线”。重新实现者在接线 prefetch 路径时,不得假定对称的 evict 路径存在。

---

## 5. 完整示例 — `impure_use_iteration_mask` → `RaggedDotExpander`

### 目的

端到端示例展示了一个标志如何驱动编译器变换:`xla_tpu_impure_use_iteration_mask`(读取**类别 (A)**,即 `FLAGS_` 锚定的直接读取,用来对照注册表中介路径)选择将 `RaggedDot` 降低为窗口化 ragged convolution 的 iteration-mask 路径。它展示了重新实现者必须复现的完整 flag→gate→pass→HLO 链条,以及 `FLAGS_` 锚定的 `AutoOr<bool>` 如何以内联方式消费,而不是按名称消费。

### 阶段 1 — 门控

`RaggedDotExpanderShouldUseIterationMask @ 0x1d6b5d60` 是门控。其反编译后的函数体很小且精确:

```c
function RaggedDotExpanderShouldUseIterationMask(ObjectView<TCE> view, TpuTopology& topo):  // 0x1d6b5d60
    if *(int*)(*(view+8)) < 3:                              // TpuVersion >= 3 gate
        return false
    if cached_word(FLAGS_..._use_iteration_mask) set:       // FlagImpl+0x58 fast path (dword_223A7E28)
        return (cached & 0x101) != 0x100                    // AUTO=ON: true unless present-and-false
    return (FlagImpl::ReadOneWord(&FLAGS_xla_tpu_impure_use_iteration_mask) & 0x101) != 0x100

& 0x101; cmp 0x100; != 测试是 autoproto-autoor-resolution.md §3 Idiom B 中的 AUTO=ON 极性:AUTO(0x000)和显式 true(0x101)都产生 true;只有显式 false(0x100)产生 false。注意,这里是直接通过 FlagImpl::ReadOneWord @ 0x21111a60FlagImpl+0x58 缓存字)从 FLAGS_ 对象读取,也就是类别 (A) 内联读取,不是 GetFieldValueIfNotDefault。这里没有按名称查找,正是因为 impure_use_iteration_mask 有会 leaFLAGS_ 对象的消费者。

两个相邻门控共享 TpuVersion>=3 前缀:ShouldUseIterationMask @ 0x1d6b5dc0(用于 LLO SpatialMajorConvolution 发射层;use_iteration_mask OR masked-fusion skipper)和 ShouldEnableMaskedFusionIterationSkipper @ 0x1d6b5d20(通过 ReadOneBool 读取的普通 bool 标志)。expander 门控只咨询 use_iteration_mask

阶段 2 — pass 接线

PostMainFusionHloOptimize @ 0x109673b3 构造该 pass,三个 impure 标志送入构造函数:

text
AddPass<RaggedDotExpander, bool, RaggedConvContractionMode, vector<string>, Target&>   ── 0x1096d2e0
  use_iteration_mask   ← RaggedDotExpanderShouldUseIterationMask        (0x109673b3 → ctor +0x08)
  contraction_mode     ← FLAGS_xla_tpu_impure_contract_ragged_conv_with (0x223a7cf8, RaggedConvContractionMode → +0x09)
  window_bounds (g,m,k,n) ← FLAGS_xla_tpu_impure_ragged_dot_window_bounds (0x223a7d58, vector<string> → +0x10/+0x20)
make_unique<RaggedDotExpander>  ── 0x1096e360(0x30 字节:vtable+0x00,bool+0x08,mode+0x09,vec+0x10/+0x20)
```text

### 阶段 3 — `RunImpl` 验证 + 窗口解析

```c
function RaggedDotExpander::RunImpl(module):                    // 0x10fae060
    for inst in MakeComputationPostOrder(module) where inst is RaggedDot:
        require inst.ragged_dot_mode != kBatch                  // CHECK @0x8644813
        validate dims: 1 contracting, 1 lhs-non-contracting,
                       1 rhs-non-contracting, 0 batch           // CHECKs @0x9e726a5/70/3b, 0x9fc1ce3
        spec = RaggedConvSpec::FromRaggedDot(inst)              // closure @0x10fb2160
        if member.use_iteration_mask == 1:                      // cmpb $1, 0x8(this) @0x10faef92/0x10faf121
            require window_bounds.size() == 4                   // CHECK @0x9b28192
            window = PipelineWindowSpec{ g, m, k, n }           // safe_strto64_base, [0..3] @0xa117494/a109010/a115712/a108547
        ExpandRaggedDot(inst, spec, member.use_iteration_mask, window, ..., target)  // 0x10fafa20

阶段 4 — iteration mask(示例的价值所在)

ExpandRaggedDot @ 0x10fafa20RaggedDot 替换为窗口化 ragged convolution 子图,并在 use_iteration_mask 时构建一个布尔 mask,用于将 dense convolution 计算出来的跨组乘积清零。mask 生成器 CreateOutputMask @ 0x10fb2900 已按字节确认:

c
function CreateOutputMask(...):                                // 0x10fb2900
    iota   = CreateIota(output_index_dim)                       // 0x10fb2d86
    lo     = CreateBroadcast(group_lower_bound, out_shape)      // 0x10fb2d52
    hi     = CreateBroadcast(group_upper_bound, out_shape)      // 0x10fb2e9a
    ge     = CreateCompare(iota, lo, kGe=5)                      // 0x10fb2de5  (iota >= group_start)
    lt     = CreateCompare(iota, hi, kLt=?)                      // 0x10fb2ec2  (iota <  group_end)
    band   = CreateBinary(kAnd=13, ge, lt)                      // 0x10fb2f9f  (in-group range)
    mask   = CreateFusion(band)                                 // 0x10fb314b  (wrap the boolean mask)
    return mask
```text

反编译显示了 `CreateIota`、两个 `CreateBroadcast`(逐组边界)、三个 `CreateCompare`(比较码 4/5/5)、一个 `CreateBinary 13`(`kAnd`)和一个 `CreateFusion`,正是 `iota∈[group_start, group_end)` 谓词。随后 `CreateConvolutionSelectFusion @ 0x10fb31e0` 应用它:

```c
function CreateConvolutionSelectFusion(...):                   // 0x10fb31e0
    conv = CreateConvolve(ragged_conv operands)                 // 0x10fb4c9c
    sel  = CreateTernary(kSelect, mask, conv, broadcast_zero)   // 0x10fb4923/0x10fb51c2
    out  = CreateReduce(sel, AddComputation)                    // 0x10fb5346 (+ 0x10fb63a0)
    return CreateFusion(out)                                    // 0x10fb4a96/0x10fb5449

语义

ragged contracting 维度变成带 g/m/k/n 窗口的 convolution spatial 维度。dense convolution 计算跨组的完整块笛卡尔积;iteration mask Select( (iota >= group_start) AND (iota < group_end), conv, 0 ) 将每个跨组乘积清零,因此只有每个 ragged 组的有效范围保留下来,而 Reduce 将被 mask 后的窗口求和为 dense 输出。设置 use_iteration_maskTpuVersion>=3 时 AUTO=ON)会选择这种 masked lowering;禁用它(显式 false)会走非 masked 路径。这就是该示例的完整重新实现契约:标志的极性和门控、pass 的成员布局,以及 mask 的 HLO 图。

怪癖 — 同一个标志由两个不同门控以相同极性读取,但作用域不同。RaggedDotExpanderShouldUseIterationMask 只咨询 use_iteration_maskShouldUseIterationMask @ 0x1d6b5dc0(LLO conv 发射层)会将它与 masked-fusion skipper 做 OR。把一个门控的行为接线到另一个上的重新实现者,会在 skipper 场景下产生分歧。该标志在两个位置都是类别 (A):直接从 FLAGS_ 内联读取,从不按名称读取。


相关组件

组件关系
GetFieldValueIfNotDefault<T> @ 0x1c6f1a80按名称反射读取器 — 注册表中介读取
FlagFieldMappings @ 0x2257ef50(构造函数 0x1d753ce0使名称可解析的 flag↔field 桥
SetFieldFromFlagString @ 0x1d73fcc0写入侧 — CLI 标志字符串 → TCE proto 字段
LegacyEvictions/PrefetchesFlagRegistry @ 0x22579860 / 0x225798a0TpuVersion 划分的字段名称注册表(规范类别 B 示例)
RaggedDotExpander @ 0x10fae060 + CreateOutputMask @ 0x10fb2900完整 flag→transform:iteration-mask ragged-conv lowering
AutoOr<bool>::FromProtoOrDie @ 0xf795300类别 (A) 消费者用来替代按名称路径的内联解析器

交叉引用