注册表中介标志
本页所有地址均适用于
libtpu-0.0.40-cp314wheel 中的libtpu.so(构建libtpu_lts_20260413_b_RC00,build-id md589edbbe81c5b328a958fe628a9f2207d,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> @ 0x1c6f1a80、FlagFieldMappings flag↔field 桥 @ 0x2257ef50、按 TpuVersion 划分的 LegacyEvictions/LegacyPrefetches 同步标志注册表、写入侧 SetFieldFromFlagString @ 0x1d73fcc0,以及一个标志驱动编译器变换的完整示例:xla_tpu_impure_use_iteration_mask → RaggedDotExpander iteration-mask 降低。
参照框架是 XLA 的 xla::DebugOptions proto 后端标志,但带有一个 libtpu 特有的细节,重新实现者必须复现。CLI 标志(absl::CommandLineFlag,即 .data 中的 FLAGS_<name> 对象)和 TCE proto 字段是运行时桥接的两个独立事物:FlagFieldMappings 的 NoDestructor 构造函数遍历每个 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_scheduler 和 explicit_evict_memory_limit_kib 都已注册(因此 --helpfull 会列出它们,CLI 也能解析它们),但没有静态可达的消费者:它们的 FLAGS_ 对象恰好只有一个代码 xref,即注册用的 lea;并且它们的名称不存在于序列化的 TCE FileDescriptorProto 中,所以按名称读取器的 FindFieldByName 会失败。它们是残留物。本页结构如下:三种读取类别;按名称反射读取器和桥;证明这两个具名标志未接线,并说明实际驱动这些行为的机制;最后是完整的 RaggedDot flag→transform 示例。
对于重新实现,契约是:
- 三种读取类别 — (A)
FLAGS_锚定的直接读取(lea FLAGS_;FlagImpl+0x58;AutoOr内联解析),(B) 按名称 TCE 反射(注册表中介路径),(C) 已注册但未接线(名称不是 proto 字段)。重新实现者必须通过消费者的读取指令区分它们,而不是通过标志名称或类型。 - 按名称读取机制 —
GetFieldValueIfNotDefault<T>=FindFieldByName→GetFieldValue反射 → 与全默认 env 求差 →optional<T>;FlagFieldMappings双向FlatHashMap桥;写入侧SetFieldFromFlagString。 - 完整 flag→transform 示例 —
impure_use_iteration_mask(读取类别 A,AutoOr<bool>,AUTO=ON,TpuVersion>=3门控)→RaggedDotExpander→Iota + Compare + AND + Select(mask, conv, 0)iteration mask,用于在 ragged convolution 中将跨组乘积清零。
| 按名称读取器 | GetFieldValueIfNotDefault<long> @ 0x1c6f1a80 — FindFieldByName + TpuCompEnvReflection::GetFieldValue + 与默认值求差 |
| Flag↔field 桥 | FlagFieldMappings::GetInstance @ 0x2257ef50 · 构造函数 @ 0x1d753ce0(双向 FlatHashMap,每字段一次 FindCommandLineFlag) |
| 写入侧 | SetFieldFromFlagString @ 0x1d73fcc0 — GetFieldForFlag + 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 @ 0x223a7dd0 — AutoOr<bool>,FlagOps 0x1d6b5840,6 个 xref(5 个消费者 + 注册器) |
| 示例变换 | RaggedDotExpander::RunImpl @ 0x10fae060 → CreateOutputMask @ 0x10fb2900 → CreateConvolutionSelectFusion @ 0x10fb31e0 |
| 置信度 | CONFIRMED(字节锚定 vs 反编译),除非某行或提示另有说明 |
1. 三种读取类别
目的
TCE 编译旋钮的值可以通过三条不同路径到达消费者。这一区分很重要,因为它会同时改变重新实现者查找消费者的位置,以及该标志是否在运行时真的有影响。对每个标志而言,这些类别互斥,并且标志所属类别只能从其消费者的读取指令中看出来。
类别
| 类别 | 消费者处的读取指令 | 示例标志 | 效果 |
|---|---|---|---|
(A) FLAGS_ 锚定的直接读取 | lea FLAGS_<name>(%rip) → FlagImpl::ReadOneWord/ReadOneBool(FlagImpl+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.md 和 flag-families.md 为大部分旋钮表面记录的路径:一段生成的访问器直接从 FLAGS_ 对象或 TCE 结构体偏移加载值。类别 (B) 是本页主题。类别 (C) 是退化的 (B):读取机制存在,但标志名称从未加入 TCE proto,因此按名称读取器永远找不到它。
陷阱 — 类别 (B) 和 (C) 在
FLAGS_层面共享同一个静态签名:恰好一个代码 xref,即注册lea。一次覆盖整个.text的 RIP-relativelea/mov扫描无法区分二者,因为二者都没有引用FLAGS_对象的消费者。判别器是序列化的 TCE proto:类别 (B) 标志的名称是 proto 字段 token(所以FindFieldByName成功);类别 (C) 标志的名称不是。只按 xref 数量分类的重新实现者会把每个未接线标志都误归为单纯的注册表中介标志。
标志值如何进入 proto(写入侧)
对于类别 (B),标志值必须先写入 TCE proto 字段,随后任何按名称读取才能观察到它。这个桥是 SetFieldFromFlagString @ 0x1d73fcc0:
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 上不同:
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::FindFieldByName | 0x20e57900 | 名称 → FieldDescriptor(null ⇒ 不是字段) |
TpuCompEnvReflection::GetFieldValue | 0x1d7523a0 | 反射读取 → 覆盖所有 arm 类型的 variant |
GetTpuCompEnvWithDefaultValues | 0x1d73f100 | 读取差分所针对的全默认 env |
FlagFieldMappings ctor | 0x1d753ce0 | 构建双向 FlatHashMap(flag↔field) |
FlagFieldMappings::GetInstance()::mappings | 0x2257ef50 | NoDestructor 单例存储(访问器被内联) |
TpuCompEnvReflection::GetFlagForField | 0x1d74ad40 | field→flag 查找(缺失时 FATAL) |
SetFieldFromFlagString | 0x1d73fcc0 | 写入侧 — 标志字符串 → 通过反射写入 proto 字段 |
RegisterCommandLineFlag | 0x21114cc0 | 注册器中每个标志的注册调用 |
3. 旧式同步标志注册表(类别 (B) 示例)
目的
最干净的存活类别 (B) 消费者是 MSA(memory-space assignment)逐出/预取限制读取器。逐出和预取的 outstanding-op 限制是按 TPU 系列划分的 TCE 字段,带有系列特定名称(Jellyfish 用 xla_jf_…,还有 xla_vf_…、xla_zf_…、xla_gf_…)。消费者不是为每个系列硬编码一个偏移,而是在按 TpuVersion 划分的注册表中查找字段名称,然后按名称读取该字段。
算法
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 @0x223c72a8 → 1 xref: lea @0x21360e44
FLAGS_xla_tpu_explicit_evict_memory_limit_kib @0x223c50d0 → 1 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 @0x223a7dd0 有 6 个 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_scheduler:movw $0 ⇒ AUTO;explicit_evict:movq $0; movb $0,0x8 ⇒ {0, has0} ⇒ AUTO)。它们是形态良好但没有读取器的标志。
| 标志 | FLAGS_ | 类型 | 唯一 xref | 实际驱动行为的机制 |
|---|---|---|---|---|
xla_tpu_enable_lem_scheduler | 0x223c72a8 | AutoOr<bool>(FlagOps 0x1d6b5840) | lea @ 0x21360e44(注册) | 运行时谓词 ModuleContainsLEMSparseCoreInstruction @ 0x13853280 |
xla_tpu_explicit_evict_memory_limit_kib | 0x223c50d0 | AutoOr<int64_t>(FlagOps 0x1d700120) | lea @ 0x21360976(注册) | 旧式逐版本 vmem-evict TCE 字段(§3) |
实际驱动 LEM scheduler 的机制
LEM-scheduler 决策不是该标志。在 RunHloScheduler @ 0x1096fac0 中:
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 @ 0x21111a60(FlagImpl+0x58 缓存字)从 FLAGS_ 对象读取,也就是类别 (A) 内联读取,不是 GetFieldValueIfNotDefault。这里没有按名称查找,正是因为 impure_use_iteration_mask 有会 lea 其 FLAGS_ 对象的消费者。
两个相邻门控共享 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 标志送入构造函数:
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 @ 0x10fafa20 将 RaggedDot 替换为窗口化 ragged convolution 子图,并在 use_iteration_mask 时构建一个布尔 mask,用于将 dense convolution 计算出来的跨组乘积清零。mask 生成器 CreateOutputMask @ 0x10fb2900 已按字节确认:
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_mask(TpuVersion>=3 时 AUTO=ON)会选择这种 masked lowering;禁用它(显式 false)会走非 masked 路径。这就是该示例的完整重新实现契约:标志的极性和门控、pass 的成员布局,以及 mask 的 HLO 图。
怪癖 — 同一个标志由两个不同门控以相同极性读取,但作用域不同。
RaggedDotExpanderShouldUseIterationMask只咨询use_iteration_mask;ShouldUseIterationMask @ 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 / 0x225798a0 | 按 TpuVersion 划分的字段名称注册表(规范类别 B 示例) |
RaggedDotExpander @ 0x10fae060 + CreateOutputMask @ 0x10fb2900 | 完整 flag→transform:iteration-mask ragged-conv lowering |
AutoOr<bool>::FromProtoOrDie @ 0xf795300 | 类别 (A) 消费者用来替代按名称路径的内联解析器 |
交叉引用
- overview.md — 三层配置流水线;本页负责 Stage 3 的按名称(注册表中介)读取变体
- autoproto-autoor-resolution.md — 偏移键控的
AutoOr模型;与此处按名称读取形成对照,并说明示例复用的 AUTO=ON 极性 - tce-field-offsets-defaults.md — 字节精确的 field#→offset→default 映射;本页按名称路径会绕过直接偏移读取
- flag-families.md — 标志族,以及示例所用的
impure_…族 - xla-flag-atlas.md — 更广泛的标志目录和命名约定