编译缓存
地址、构建 ID 和符号名称适用于
libtpu-0.0.40-cp314轮中的libtpu.so(构建 ID89edbbe81c5b328a958fe628a9f2207d)。其他版本有所不同;将每个 VA 视为版本固定。
摘要
编译 TPU 程序的成本很高——完整的 HLO→MLIR→LLO 下降、MSA、调度和捆绑打包运行端到端(请参阅 编译阶段.md)。因此,如果 libtpu 可以将请求识别为相同的,则它永远不会重新编译它已经构建的程序。识别是一个缓存键问题:运行时必须将整个编译请求(程序文本、编译器标志和物理目标)减少为一个简短的可比较令牌,查找该令牌,并在未命中时编译一次并将结果存储在其下。此页面拥有该机制:**哈希到密钥中的内容、查找/存储路径以及缓存条目生命周期。**最终在命中条目内部的序列化可执行文件(TpuExecutable / TpuProgram 有线格式)由 tpu-program-serialization.md 所有,并在此处被视为不透明有效负载。
有三个不同的缓存层,将它们合并是第一个重新实现陷阱。最外面的是 XLA JIT 缓存,tensorflow::DeviceCompilationCache<T>(来自 device_compilation_cache.h):由 DeviceCompilationClusterSignature 键控的进程内 flat_hash_map(函数名称加上规范化参数形状/值),决定 XLA 集群是否需要(重新)编译。它下面是 TPU C-API 密钥,tensorflow::tpu::TpuCompilationCacheKey,由 TpuCompile_CreateCompilationCacheKey (0xf6a2080) 在 tpu_util_c_api.cc 中构建:Fingerprint2011 对由 13 个命名输入组成的 StrCat 组装字符串(HLO/MLIR 指纹、副本计数、拓扑)进行的摘要边界、保证常量大小、形状前缀、嵌入分区指纹)加上条件设备分配尾部。最里面的是 TFRT 组缓存,tensorflow::tfrt_tpu::TpuCompilationCache (tpu_compilation_cache.cc):该层实际上“保存”已编译的程序,通过三层查找(内存中 → 协调服务 → 磁盘上持久缓存)解析密钥,并管理条目引用计数和逐出。
页面按层组织,最外层在前。每一层都具有相同的 ### Purpose / ### Cache Key / ### Lookup & Store / ### Algorithm / ### Function Map 语法,因此读者可以并排比较这三个密钥。关闭单元涵盖持久磁盘路径格式和条目生命周期(引用计数、逐出、从序列化恢复)。
重新实现,合约为:
- TPU 缓存密钥是
Fingerprint2011,位于由十三个命名字段加上条件设备分配尾部组成的确定性有序字符串。重现字段集、串联顺序(与dump_vars标签顺序不同 -shapes_prefix最后附加)、分隔符(:和,)和两阶段哈希(首先对嵌入分区原型进行指纹识别,然后对整个前缀进行指纹识别)。任何省略或重新排序的字段都会默默地更改密钥并导致虚假丢失。 - XLA 端密钥是签名,而不是哈希 -
DeviceCompilationClusterSignature存储函数名称加上每个参数的规范形式(值携带的常量张量,(dtype, shape)携带的非常量),与自定义Hash/==进行比较。它控制是否TpuCompile_CreateCompilationCacheKey甚至被调用。 - 查找是三层的。
TpuCompilationCache::LookUpInternal(0xf7a72c0) 首先尝试内存中的映射,然后是协调服务,然后是持久磁盘缓存;每一层的未命中是触发真正编译的唯一路径。 - 条目是引用计数且可LRU可清除的。 缓存保存具有显式
Release/DiscardEntryRef/MarkOldestEntryForEviction的CompiledSubgraph/TpuCompilationCacheEntry对象;正在执行的程序固定其条目。 - 持久路径是文件系统前缀,而不是哈希目录:
JoinPath(dir, "CL" + guaranteed_const_fp + "_" + base_key),由ConstructCacheEntryFilepathPrefix(0xe8cab00)构建。
| TPU 密钥生成器 | TpuCompile_CreateCompilationCacheKey — 0xf6a2080 (tpu_util_c_api.cc:149-193) |
| 关键摘要 | Fingerprint2011 (0x20d6cd40) → 十进制字符串;两阶段 |
| XLA 签名 | tensorflow::DeviceCompilationClusterSignature::Build — 0xfb01f40 |
| XLA 内存缓存 | DeviceCompilationCache<T>::{LookupOrCreate,Store} — 0xe9932e0 / 0xe994a60 |
| TFRT组缓存 | tensorflow::tfrt_tpu::TpuCompilationCache (tpu_compilation_cache.cc) |
| 三层查找 | TpuCompilationCache::LookUpInternal — 0xf7a72c0 |
| 持久路径 | ConstructCacheEntryFilepathPrefix — 0xe8cab00 ("CL" + fp + "_" + key) |
| 缓存条目类型 | tfrt::tpu::TpuCompilationCacheEntry(构造函数0xf7bd100) |
| 源文件 | tpu_util_c_api.cc、device_compilation_cache.h、tpu_compilation_cache.cc |
第 1 层 — XLA 集群签名 (DeviceCompilationCache)
用途
这是最外面的门,也是 JAX/XLA 调用者首先点击的门。当即将编译 XLA 集群(从 NameAttrList 及其参数列表派生的 HLO 计算)时,DeviceCompilationCache<T> 决定在此过程中是否已经编译了等效集群。它以可执行类型为模板 - xla::PjRtLoadedExecutable (0xe9932e0) 和 xla::LocalExecutable (0xe986a20) 实例化都存在于二进制文件中,共享相同的 LookupOrCreate/Store 代码。它是来自 third_party/tensorflow/compiler/jit/device_compilation_cache.h 的纯内存缓存(Store 日志站点引用了 device_compilation_cache.h:259);它没有自己的持久性,并将实际编译委托给 DeviceCompiler<T>::CompileStrict/CompileAsynchronous。
缓存密钥
密钥是 tensorflow::DeviceCompilationClusterSignature,由 DeviceCompilationClusterSignature::Build (0xfb01f40) 构建(未散列)为结构化值:
// DeviceCompilationClusterSignature::Build (sub_0xfb01f40)
// canonical_function : a DeviceCompilationCanonicalFunction (name + attrs)
// args : Span<const XlaArgument>
function Build(canonical_function, args):
sig.name = canonical_function.name // copied string, +0..+15
AppendArguments(sig, args) // sub_0xfb01a80
return sig
```text
`AppendArguments` (0xfb01a80) 遍历每个 `XlaArgument` 并附加一个 `std::variant`:**编译时常量**参数被**按值**捕获为 `tensorflow::Tensor`(因此具有不同常量值的两个调用获得不同的签名),而非常量参数被捕获为 `pair<DataType, InlinedVector<int64,4>>` — 仅其数据类型和形状。签名通过 `DeviceCompilationClusterSignature::Hash` 和相等性进行比较,这就是为什么下面的映射是 `FlatHashMapPolicy<DeviceCompilationClusterSignature, unique_ptr<Entry>>, Hash>` 而不是字符串键控映射的原因。
> **明白了 —** 签名*对常量值敏感,对变量形状敏感*。仅形状上的键的重新实现将在常量参数不同的调用之间错误地共享条目(例如,`slice` 起始索引折叠为保证常量)。捕获常量的完整 `Tensor` 是经过深思熟虑的,并且加载顺序很重要:`AppendArguments` 保留参数顺序。
### 查找和存储
`LookupOrCreate` (0xe9932e0) 和 `Store` (0xe994a60) 都在单个 `absl::Mutex` 下运行,守护着地图。仅当签名不存在时,`LookupOrCreate` 才会放置一个新的空 `Entry`(已分配 `operator new(0x30)`,48 字节),然后返回条目当前状态的快照。 `Store` 通过签名重新查找条目,并写回编译器生成的四个可选字段中的任意一个。
```c
// Entry layout (48 bytes, operator new(0x30); offsets from the Entry base, byte-confirmed in Store)
// +0 : absl::Mutex mu (per-entry lock; Store locks Entry+0)
// +8 : DeviceCompileState compile_state (DWORD; kUncompiled / kCompiling / kCompiled)
// +16 : int64 request_count (bumped on each LookupOrCreate; init 0)
// +24 : absl::Status compilation_status (StatusRep*, refcounted; init 1 = inline OK)
// +32 : XlaCompilationResult* compilation_result (owned; deleted on overwrite)
// +40 : T* executable (owned; vtable-deleted on overwrite)// DeviceCompilationCache<T>::LookupOrCreate (sub_0xe9932e0)
function LookupOrCreate(signature) -> EntrySnapshot:
lock(mutex_)
entry = map_.try_emplace(signature, new Entry{}) // EmplaceDecomposable
entry.request_count += 1 // +16
snapshot.compile_state = entry.compile_state // +8
snapshot.status = entry.compilation_status // +24, Ref()'d
snapshot.request_count = entry.request_count
snapshot.{result,exe} = entry.{result,exe} // +32,+40 (aliased, not owned)
unlock(mutex_)
return snapshot
// DeviceCompilationCache<T>::Store (sub_0xe994a60)
function Store(signature, opt_state, opt_status, opt_result, opt_exe):
lock(mutex_)
entry = map_.try_emplace(signature, new Entry{}) // re-find or create
unlock(mutex_)
lock(entry.mutex) // per-entry mutex at Entry+0
if opt_state.has_value(): entry.compile_state = *opt_state
if opt_status.has_value(): entry.compilation_status = *opt_status // swap+Unref old
if opt_result.has_value(): delete entry.result; entry.result = release(*opt_result)
if opt_exe.has_value(): vtable_delete entry.exe; entry.exe = release(*opt_exe)
unlock(entry.mutex)
VLOG(4) << "Added/updated cache entry: key=" << sig.HumanString()
<< ", entry=" << entry.DebugString()
```text
> **QUIRK —** `Store` 采用四个 `std::optional` 字段并仅更新存在的字段。编译器首先存储 `compile_state = kCompiling`(声明插槽),然后才存储 `compilation_result` + `executable` + `compile_state = kCompiled`。观察 `kCompiling` 的并发 `LookupOrCreate` 预计会等待未来状态而不是重新编译 - 缓存本身不会阻塞,它会显示状态。
### 功能图
| 功能 | 地址 | 角色 |
|---|---|---|
| `DeviceCompilationClusterSignature::Build` | `0xfb01f40` | 构建名称+args签名 |
| `AppendArguments` | `0xfb01a80` | 附加每个参数变量(const→Tensor,var→dtype/shape) |
| `DeviceCompilationCache<PjRt…>::LookupOrCreate` | `0xe9932e0` | Emplace + 互斥锁下的快照 |
| `DeviceCompilationCache<PjRt…>::Store` | `0xe994a60` | 回写可选字段 |
| `DeviceCompilationCache<LocalExe>::LookupOrCreate` | `0xe986a20` | 相同,`LocalExecutable` 实例化 |
| `DeviceCompilationCache<LocalExe>::Store` | `0xe988d00` | 相同,`LocalExecutable` 实例化 |
| `DeviceCompiler<…>::CompileStrict` | `0xe993860` / `0xe986fa0` | 缺少路径:进行真正的编译 |
| `DeviceCompiler<…>::CompileAsynchronous` | `0xe993420` / `0xe986b20` | 异步未命中路径 |
---
## 第 2 层 — TPU 编译缓存密钥
### 用途
当第 1 层未命中并且编译器到达 TPU 后端时,请求将简化为 `tensorflow::tpu::TpuCompilationCacheKey` — TFRT 组缓存(第 3 层)和持久磁盘缓存在该令牌下建立索引。单个构建器是 `TpuCompile_CreateCompilationCacheKey` (0xf6a2080),通过 `learning/45eac/tfrc/executor/stream_executor/tpu_util_c_api.cc` 中的 C-ABI 填充层导出(`LOG(FATAL)` 和 `VLOG` 站点引用第 149、152、178、193 行)。它是该页面上最重要的函数,因为它的字段集*是*缓存的程序标识概念。
### 缓存密钥
密钥是一个十进制字符串:`StrCat` 组装的 **前缀** 字符串的 `Fingerprint2011`,其本身是该 64 位指纹的 `to_string`。构建器为其 `VLOG` 转储(`tpu_util_c_api.cc`、`$_1` lambda、`dump_vars<13ul, 13ul, …>`)携带一个 13 条目 `dump_vars` 标签表;这十三个标签命名程序标识输入,第十四个标签 (`device_assignment`) 有条件地附加,但 **不在** 中。下表按名称列出了输入; **##** 列是转储表标签顺序,它与连接顺序*不*相同(请参阅伪代码后面的注释)。
| # | 字段(dump_vars 标签) | 源码 |
|---|---|---|
| 1 | `function_name` | `std::string(property.function_name)` |
| 2 | `function_library_fingerprint` | `property.function_library_fingerprint` |
| 3 | `mlir_module_fingerprint` | `property.mlir_module_fingerprint` |
| 4 | `num_replicas` | `property.num_replicas` |
| 5–7 | `chip_bounds.{x,y,z}` | `topology.chip_bounds().{x,y,z}`(拓扑+0x58区域) |
| 8–10 | `wrap.{x,y,z}` | `topology.wrap().{x,y,z}`(拓扑+0xA0/+0xA2) |
| 11 | `shapes_prefix` | `std::string(property.shapes_prefix)` |
| 12 | `guaranteed_constants_size` | `property.guaranteed_constants_size` |
| 13 | `embedding_partitions_fingerprint` | `Fingerprint2011(embedding_partitions_proto)` |
| — | `device_assignment`(有条件,不在转储表中) | `:device_assignment:` + 加入的 id,**或** `:default_device_assignment` |
```c
// TpuCompile_CreateCompilationCacheKey (sub_0xf6a2080, tpu_util_c_api.cc:149)
// property : compilation request (function name/lib fp, mlir fp, replicas, shapes, …)
// mesh : TpuMeshCommonState* (carries topology + embedding partitions)
function CreateCompilationCacheKey(property, mesh) -> TpuCompilationCacheKey:
CHECK(mesh != nullptr) // line 149
// (a) fingerprint the embedding-partitions proto first
eb_str = mesh.embedding_partitions_proto().SerializeToString() // CHECK ok, line 152
eb_fp = Fingerprint2011(eb_str)
topo = mesh.tpu_topology()
// (b) assemble the ordered prefix string
// (StrCat<...,int,...,int,...,int,...,int,...,bool,...,bool,...,bool,
// ...,char const*,...,string> in the decompile)
prefix = StrCat(property.function_name, // field 1 (string, v103/a10)
":", property.function_library_fingerprint, // 2 (FastIntToBuffer)
":", property.mlir_module_fingerprint, // 3 (FastIntToBuffer)
":", property.num_replicas, // 4 (int, a16)
":", topo.chip_bounds.x, ",", .y, ",", .z, // 5-7 (3×int, topo+0x58 region)
",", topo.wrap.x, ",", .y, ",", .z, // 8-10 (3×bool, topo+0xA0/+0xA2)
":", property.guaranteed_constants_size, // 12 (int, a9)
":", to_string(eb_fp)) // 13 (embedding fp, v87)
// (c) device-assignment tail (field 14): explicit ids vs. default
if device_grid_matches_replicas(property): // line ~220 shape check
if property.device_assignment == nullptr:
StrAppend(prefix, ":default_device_assignment")
else:
StrAppend(prefix, ":device_assignment:", Join(ids, ","))
StrAppend(prefix, property.shapes_prefix) // field 11 (string, a8) — appended last
VLOG(1) << "prefix_raw = " << prefix // line 178
// (d) the key proper is the fingerprint of the whole prefix
key.fingerprint = to_string(Fingerprint2011(prefix))
key.prefix = prefix
VLOG(1) << "CompilationCacheKey:" << key // line 193
return key // {fingerprint, prefix} pair连接顺序 —
dump_vars标签表是构建器的显示顺序,而不是连接顺序。从反编译来看,主StrCat运行function_name,然后:-加入了function_library_fingerprint、mlir_module_fingerprint、num_replicas,然后,-加入了chip_bounds.{x,y,z}和wrap.{x,y,z},然后:-加入guaranteed_constants_size和嵌入-分区指纹。设备分配尾部接下来是StrAppended,shapes_prefix是StrAppended 最后(它是设备分配块之后的最后一个StrAppend)。重新实现必须匹配此串联顺序,而不是转储标签顺序,否则指纹会发散。各个
chip_bounds/wrap组件的精确 int-vs-bool 渲染(反编译的StrCat模板实例化四个int和三个boolAlphaNum 插槽)是 Hex-Rays 工件,未固定每个组件;字段 names 和上面的总体顺序以字节方式锚定到dump_vars标签和StrCat/StrAppend控制流。
保证常数贡献本身就是一个指纹,由 TpuCompile_CreateGuaranteedConstFingerprint (0xf6a2040) 单独计算并与 FingerprintCat2011 结合:
// TpuCompile_CreateGuaranteedConstFingerprint (sub_0xf6a2040)
function CreateGuaranteedConstFingerprint(running_fp, const_bytes) -> uint64:
return FingerprintCat2011(running_fp, Fingerprint2011(const_bytes))
```text
> **GOTCHA —** 密钥折叠 *两者* 内容指纹(HLO/MLIR 模块 fp,字段 3;常量大小,字段 12;嵌入分区 fp,字段 13)*和* 物理目标(拓扑芯片边界 + 包装,字段 5-10;设备分配尾部)。针对不同拓扑或设备分配编译的两个字节相同的程序获得不同的密钥 - 正确,因为编译的 `TpuProgram` 是特定于目标的。仅在程序文本上键入键的重新实现将为 v5e-2x2 请求提供 v4-2x2x1 二进制文件,并在加载时崩溃。
>
> **QUIRK —** 设备分配尾部通过反编译行 220 处的检查进行门控。当 `(replicas.lo × replicas.hi == num_cores) || (replicas.lo ∉ {1, num_cores})` 时,尾部被 **附加**,其中 `num_cores` 是从 `topology+0x80` 读取的,两个 `replicas` 半部来自 64 位副本字段;否则,显式/默认分配字段将被**完全省略**,而不是默认的 - 因此密钥长度是可变的。准确地再现谓词,否则同一程序的单副本和多副本运行将意外地发生冲突或分歧。 (这里的网格/副本字段标识是从守卫的算术形状推断出来的,而不是从标记的符号中推断出来的——对哪个字段是哪个字段的置信度中等;谓词结构本身是字节锚定的。)
### 功能图
| 功能 | 地址 | 角色 |
|---|---|---|
| `TpuCompile_CreateCompilationCacheKey` | `0xf6a2080` | 组装前缀+指纹→密钥 |
| `TpuCompile_CreateGuaranteedConstFingerprint` | `0xf6a2040` | `FingerprintCat2011(fp, Fingerprint2011(const))` |
| `TpuCompile_DestroyCompilationCacheKey` | `0xf6a2e60` | 释放关键结构 |
| `Fingerprint2011` | `0x20d6cd40` | `string_view` 的 64 位非加密摘要 |
| `FingerprintCat2011` | `0x20d6d0e0` | 组合两个64位指纹 |
| `TpuProgram_GetFingerprint` | `0xe8bed60` | 从编译的 `TpuProgram` 读取指纹(+160→+648) |
| `TpuExecutable_Fingerprint` | `0xeabea40` | 序列化可执行文件的指纹 |
> **注意 —** `Fingerprint2011` 是 Farmhash 派生的 TF 指纹,*不是*加密哈希。缓存信任它的身份;原则上,对手可以强制发生冲突,但威胁模型是避免重新编译,而不是完整性。 `TpuProgram_GetFingerprint` 读取存储在已编译程序*内部*的指纹(`program+160 → +648` 中的内联字符串优化字段),用于根据加载时的密钥验证反序列化程序。
---
## 第 3 层 — TFRT 组缓存 (`tfrt_tpu::TpuCompilationCache`)
### 用途
这是实际“存储已编译程序”并在运行时决定编译与重用的层。第 2 层键上的 `tensorflow::tfrt_tpu::TpuCompilationCache` (`learning/45eac/tfrt/tf_tpu/tpu_compilation_cache.cc`) 键,保存按程序组分组的 `tfrt::tpu::TpuCompilationCacheEntry` 对象,并按优先级顺序通过三层解析查找。 “组”是一组一起插入和驱逐的相关条目(例如,一个分片计算的每个核心程序); `EmplaceGroupEntry`/`RemoveGroup`/`RemoveGroupEntries` 以组粒度运行,而各个条目则带有自己的引用计数。
### 缓存密钥
第 3 层在第 2 层指纹字符串上建立索引,但首先使用 `ConvertToBaseKey` (0xf7bf760) 对其进行规范化,这会去除每个核心的后缀,以便组中的所有核心共享一个基本密钥。持久磁盘索引使用进一步派生的路径前缀(请参阅下面的*持久磁盘缓存*)。
### 查找和存储
命中路径是 `LookUpInternal` (0xf7a72c0),一个三层解析器:
```c
// TpuCompilationCache::LookUpInternal (sub_0xf7a72c0, tpu_compilation_cache.cc:~1097)
function LookUpInternal(base_key) -> AsyncValueRef<Entry>:
// TIER 1 — in-memory map (shared lock; walk tombstone chain)
lock_shared(cache_mutex_)
node = entries_head_; while node.flags & 3: node = node.next // skip dead slots
result = GetTpuCompilationCacheEntry(node) // hit if present
unlock_shared(cache_mutex_)
if result.is_present: return result
// TIER 2 — coordination service (multi-host: peer may already hold it)
if coord_service_ && coord_service_.is_available():
return LookUpWithCoordService(base_key)
// TIER 3 — persistent on-disk cache
if persistent_cache_enabled_: // entries_[20]
st = LoadGroupFromPersistentCache() // sub: read RecordReader
if st == kOk:
node = entries_head_; while node.flags & 3: node = node.next
return GetTpuCompilationCacheEntry(node) // now in-memory
else:
return MakeErrorAsyncValueRef(st) // propagate load error
// total miss with no persistence → hard error (the compile path runs elsewhere)
return MakeErrorAsyncValueRef(
Internal("Cannot access the coordination service or persistent "
"cache to fetch the compilation cache entry.")) // line 1097存储/编译路径为 CompileGroup (0xf7a4f60) → MaybeCompileGroup (0xf7a6200) → EmplaceGroupEntry (0xf7a3a00)。 MaybeCompileGroup 是重复数据删除点:它在启动编译之前在独占锁下重新检查缓存,因此对同一密钥的两个并发首次请求会编译一次并共享结果。 RestoreFromSerialized (0xf7a2ca0) 是持久缓存加载器的对应项 - 它从序列化 blob (TpuCompilationEntrySource) 重建条目而不是编译。
QUIRK — 内存层遍历侵入节点链(+16 处的
node = node.next),并跳过设置了+4 处低两个标志位的任何节点。这些位标记了逻辑删除/被驱逐的条目;简单的flat_hash_map::find重新实现将返回一个正在驱逐的条目。共享锁遍历加上标志检查是正确的读取协议。
功能图
| 功能 | 地址 | 角色 |
|---|---|---|
TpuCompilationCache::LookUpInternal | 0xf7a72c0 | 三层查找(mem→coord→disk) |
TpuCompilationCache::LookUpWithCoordService | 0xf7a74a0 | Tier-2 多主机查找 |
TpuCompilationCache::LoadGroupFromPersistentCache | 0xf7a7de0 | Tier-3 磁盘读取 |
TpuCompilationCache::CompileGroup | 0xf7a4f60 | 缺少路径:编译组 |
TpuCompilationCache::MaybeCompileGroup | 0xf7a6200 | 锁下重新检查,dedup编译 |
TpuCompilationCache::EmplaceGroupEntry | 0xf7a3a00 | 将编译好的组插入到地图中 |
TpuCompilationCache::GetEntriesFrom | 0xf7a1f80 | 解析组 → 每个核心条目 |
TpuCompilationCache::RemoveGroup | 0x21389f60 | 删群 |
TpuCompilationCache::RemoveGroupEntries | 0xf7a85c0 | 删除一组中的 N 个条目 |
TpuCompilationCache::RestoreFromSerialized | 0xf7a2ca0 | 从序列化 blob 重建条目 |
TpuCompilationCache::ConvertToBaseKey | 0xf7bf760 | 剥离每核后缀 → 组基键 |
tfrt::tpu::TpuCompilationCacheEntry::Program ctor | 0xf7bc500 | 构建条目的 Program 有效负载 |
tfrt::tpu::TpuCompilationCacheEntry ctor | 0xf7bd100 | 构建完整的缓存条目 |
持久磁盘缓存
持久缓存将 TpuCompilationCacheKey 转换为 文件系统路径前缀,而不是哈希桶目录。 ConstructCacheEntryFilepathPrefix(0xe8cab00,在匿名命名空间中)构建它:
// ConstructCacheEntryFilepathPrefix (sub_0xe8cab00)
// dir : persistent cache directory (string_view)
// key : TpuCompilationCacheKey
// topology : tpu::TpuTopology
// guaranteed_fp : int64 guaranteed-const fingerprint
function ConstructCacheEntryFilepathPrefix(dir, key, topology, guaranteed_fp) -> path:
base = ConstructPersistentCompilationCacheKey(key, topology) // derived stable key
file = JoinStrings({ "CL", to_string(guaranteed_fp), base }, "_") // "CL<fp>_<base>"
return JoinPath(dir, file) // tsl::io::JoinPathImpl
```text
磁盘记录本身由 `tfrt_tpu::RecordReader`(ctor 0xf7beda0)和 `RecordWriter`(ctor 0xf7be740)读取/写入 — TFRecord 样式的框架 blob 格式也用于序列化的 `TpuProgram` 有效负载。缓存目录及其只读/读写模式由absl标志控制(二进制文件带有`tfrt_tpu_pjrt_client_*`标志系列);持久目录的确切标志名称没有单独跟踪(标志身份的置信度低),但上面的路径前缀构造是字节确认的。
> **注意 —** `"CL"` 文字是持久缓存记录类型标记(该文件是*编译的*程序)。保证常量指纹被提升到*文件名*中——与键内字段 12(大小)分开——因此常量*值*(不仅仅是它们的字节大小)的变化会产生一个不同的磁盘文件,即使内存中的键看起来很稳定。这是第 1 层按值捕获常量 `Tensor`s 的磁盘上模拟。
---
## 入门生命周期
缓存条目在其整个生命周期内进行引用计数:在未命中时创建,在程序执行时固定,并在 LRU 压力下逐出。 TFRT 层的生命周期原语位于 `TpuCompilationCacheInterface` 上:
```text
LookUpInternal miss
│
▼
CompileGroup ──► MaybeCompileGroup (dedup) ──► compile ──► EmplaceGroupEntry
│ │
│ (or) RestoreFromSerialized / LoadGroupFromPersistentCache
▼ ▼
InsertEntry ───────────────────────────────────────► entry in map (refcount=1)
│
▼
execution pins entry ──► DoTpuExecute holds AsyncValueRef<TpuCompilationCacheEntry>
│
▼
Release(uid) / DiscardEntryRef ──► refcount-- ; if 0 and LRU-oldest:
│
▼
MarkOldestEntryForEviction ──► flag node (+4 bits) ; RemoveEntry frees payload| 操作 | 地址 | 角色 |
|---|---|---|
TpuCompilationCacheInterface::InsertEntry | 0xeaacde0 | 添加一个 CompiledSubgraph 到缓存 |
TpuCompilationCacheInterface::RemoveEntry | 0xeaac840 | 用钥匙删除 |
TpuCompilationCacheInterface::Release | 0xeaabde0 | 通过 uid 删除保留的引用 |
TpuCompilationCacheInterface::DiscardEntryRef | 0xeaac500 | 递减条目引用计数 |
TpuCompilationCacheInterface::MarkOldestEntryForEviction | 0xeaac340 | LRU 驱逐候选者 |
TpuCompilationCacheExternal::InitializeEntry | 0xe977c40 | 通过编译回调填充条目 |
DoTpuExecute(消耗条目) | 0xe6fa5a0 | 在运行期间保持 AsyncValueRef<TpuCompilationCacheEntry> |
tfrt::tpu::TpuCompilationCacheEntry(ctor 0xf7bd100)包装了 AsyncValueRef<Program> 以及执行器所需的元数据,而无需触及程序字节:HloInputOutputAliasConfig、输入/输出 xla::Shape 向量、CompilerMetadata、 HostTransferProto 列表、可选的 DeviceAssignment、FDO 配置和 PjrtExecutableContext。 Program 本身(ctor 0xf7bc500)是 shared_ptr<const tpu::TpuCoreProgram> 有效负载 - tpu-program-serialization.md 记录的对象。
明白了 — 该条目将程序保存为
shared_ptr<const TpuCoreProgram>,而DoTpuExecute将AsyncValueRef带到条目,而不是程序。因此,逐出不能释放正在执行中的程序:即使在MarkOldestEntryForEviction对其进行标记后,执行引用也会使条目(以及传递性的程序)保持活动状态。在驱逐标志上而不是在最后引用上释放程序的重新实现将在并发执行+驱逐下使用释放后。
各层如何组成
单个 PJRT_Client_Compile 调用(PJRT 入口点;请参阅 客户端和设备.md 和 ext-compile-phasecompile.md 中的分阶段编译)线程穿过所有三层:
PJRT_Client_Compile(program, options)
│
├─ Layer 1: DeviceCompilationClusterSignature::Build(fn, args) // 0xfb01f40
│ LookupOrCreate(sig) ─ hit? ─► return cached executable
│ │ miss → compile_state=kCompiling, fall through
│ ▼
├─ Layer 2: TpuCompile_CreateCompilationCacheKey(property, mesh) // 0xf6a2080
│ key = to_string(Fingerprint2011(prefix_of_13_fields + device_assignment))
│ ▼
├─ Layer 3: TpuCompilationCache::LookUpInternal(ConvertToBaseKey(key)) // 0xf7a72c0
│ tier1 in-mem ─ tier2 coord-svc ─ tier3 persistent("CL<fp>_<base>")
│ │ all miss
│ ▼
│ CompileGroup → MaybeCompileGroup (re-dedup) → compile → EmplaceGroupEntry
│ (optionally write persistent record via RecordWriter)
│ ▼
└─ Layer 1: Store(sig, kCompiled, status, compilation_result, executable) // 0xe994a60
return PjRtLoadedExecutable
```text
冗余是有意为之的:第 1 层以低廉的成本捕获*相同 XLA 集群*的重新编译(无指纹),第 2 层生成*目标特定*身份,第 3 层通过协调服务和磁盘捕获跨进程和交叉运行重用。每一层的未命中就是下一层的查找。
---
## 相关组件
| 名称 | 关系 |
|---|---|
| `TpuExecutable` / `TpuProgram` 序列化 | 存储在命中条目中的不透明有效负载;有线格式 |
| `DeviceCompiler<T>::CompileStrict` | 未命中路径第 1 层将真正的编译委托给 |
| `Fingerprint2011` / `FingerprintCat2011` | 每个密钥摘要所基于的哈希原语 |
| 协调服务 | `LookUpWithCoordService` 中的第 2 层多主机缓存共享 |
| PJRT客户端编译 | 驱动所有三层的公共入口点 |
## 交叉引用
- [Tpu程序序列化](tpu-program-serialization.md) — 存储在缓存命中内的 `TpuExecutable`/`TpuProgram` 有线格式;该页面将其视为不透明的有效负载
- [编译阶段](compile-phases.md) — 昂贵的端到端管道,缓存命中可让您跳过
- [编译器概述](overview.md) — 缓存程序的 IR 层堆栈的输出
- [定制呼叫降低](custom-call-lowering.md) — `tpu_custom_call` 带有自己的 `MosaicMlirCacheEntry` 子缓存,与程序缓存不同
- [PJRT客户端和设备](../pjrt/client-and-device.md) — 咨询第 1 层的 `PJRT_Client_Compile` 入口点
- [PJRT 分阶段编译扩展](../pjrt/ext-compile-phasecompile.md) — 在阶段边界之间重新序列化的多阶段编译 API
- [可执行文件执行](../pjrt/executable-execution.md) — `DoTpuExecute` 在运行期间固定缓存条目