Skip to content

TpuProgram 序列化

本页中的所有地址、字段标签和结构体偏移都适用于 libtpu-0.0.40-cp314 wheel 中的 libtpu.so(build-id 89edbbe81c5b328a958fe628a9f2207d)。其他版本会有所不同。

摘要

一个已编译的 TPU 可执行文件并不是以一个扁平 protobuf 的形式离开进程。xla::TpuExecutable::SerializeExecutable (0xf8a9300) 会组装一个 xla::TpuExecutableProto 容器,其中包含 HLO 模块 proto、逐核 LLO bundle、编译器元数据、主机传输描述符、原始编译选项,以及目标/拓扑块;随后另一个独立的分帧辅助函数 xla::TpuExecutableProtoToString (0xf8a8880) 会把该容器写入一个字节串,形式是一串长度分隔的片段,而不是单个消息。这个拆分是为了规避 protobuf 的 2 GiB 单消息上限:LLO bundle 和 HLO proto 会各自作为独立的分隔帧发出,从父消息中清除,最后再序列化只携带剩余内容的父消息。这就是 PJRT 的 Executable::Serialize 往返使用的线缆格式,也是 AOT 编译器路径(DeepseaCompiler::DeserializeExecutable, 0xfaa2660)和 C 运行时 shim(TpuExecutable_Serialize/_Deserialize)读回的格式。

这里叠放了两层容器 schema。面向 PJRT 的 xla::TpuExecutableProto 是外层信封。在其内部,字段 1 是一个 xla::DeepseaExecutableProto,即由 DeepseaExecutable::ToProto (0x134282e0) 生成的编译器内部 AOT 容器;它自身又嵌套了逐核 tpu::TpuCoreProgramProto(LLO/ISA 程序 bundle)和 xdb::CompilerMetadata。运行时 C-API 还暴露第三条更窄的路径:TpuProgram_SerializeTpuExecutable (0xe8be720) 会把单个 TpuCoreProgram 打包进 tensorflow::tpu::GetTpuProgramResponseExternal_Blob,这是用于跨缓存 RPC 边界传送单个核心程序的形式,不同于完整的 PJRT serialize。

本页记录外层信封、嵌套的 AOT 容器、逐核程序 bundle、分隔片段分帧,以及版本/兼容性戳。缓存(fingerprint 如何计算、并在调用序列化之前如何查找)见 compilation-cache.md;本页只负责线缆格式。重新实现契约如下:

  • 外层信封 xla::TpuExecutableProto:六个已填充字段、它们的标签号、存在性(has_bits)语义,以及每个字段复制自的源结构体偏移。
  • 嵌套容器DeepseaExecutableProto(LLO bundle + metadata)和 TpuCoreProgramProto(逐核 ISA 程序,一个 TensorCore/BarnaCore/SparseCore oneof)。
  • 分帧TpuExecutableProtoToString 写出的分隔式多片段字符串布局,产生“reduced proto”的字段清除过程,以及规避 2 GiB 的理由。
  • 版本/兼容性戳deepsea_version / deepsea_variant / deepsea_chip_config_name 三个 flag,以及把序列化程序固定到某个芯片代际的 TpuConfiguredProperties/TpuTopologyArgs 目标块。
外层序列化器xla::TpuExecutable::SerializeExecutable0xf8a9300
分帧辅助函数xla::TpuExecutableProtoToString0xf8a8880
外层 proto 序列化器xla::TpuExecutableProto::_InternalSerialize0xf8d0c00
AOT 容器xla::DeepseaExecutableProto(信封字段 1);ToProto 0x134282e0FromProto 0x134283e0
逐核 bundletpu::TpuCoreProgramProto — 序列化器 0x1e82ea40
AOT 反序列化xla::jellyfish::DeepseaCompiler::DeserializeExecutable0xfaa2660
C-API 序列化TpuProgram_SerializeTpuExecutable 0xe8be720TpuExecutable_Serialize 0xeabea80
C-API 反序列化TpuExecutable_Deserialize 0xeabede0TpuProgram_DeserializeFromGetTpuProgramResponseProto 0xe8be960
源 TUlearning/45eac/research/pjrt/tpu_pjrt_executable.cc
版本 flagFLAGS_deepsea_versionFLAGS_deepsea_variantFLAGS_deepsea_chip_config_nametpu_version_flag.cc

外层信封 — xla::TpuExecutableProto

目的

TpuExecutableProto 是 PJRT 级别的序列化可执行文件:它包含运行时在不重新运行编译器的情况下重新加载并入队一个程序所需的一切。它在 SerializeExecutable 内部从活动的 xla::TpuExecutable 对象逐字段构建,然后交给分帧辅助函数。它也是 TpuExecutable_Deserialize (0xeabede0) 在回程中解析的消息,不过请注意,反序列化 C-shim 解析的是内部容器 DeepseaExecutableProto,而不是完整信封(见分帧)。

字段映射

字段号和存在位直接从 TpuExecutableProto::_InternalSerialize (0xf8d0c00) 读取;源结构体偏移从 SerializeExecutable (0xf8a9300) 读取。has_bits 字节位于对象偏移 +16

字段标签类型has_bit填充来源(a2 = TpuExecutable
deepsea_executable1DeepseaExecutableProto(消息)0x08嵌套:TpuCoreProgram @ +2600/+1374 + CompilerMetadata @ +2584
hlo_module_with_config2HloModuleProtoWithConfig(消息)0x10HloModule::ToProtoWithConfigHloModule* @ +2560
host_transfers3repeated HostTransferProto0x01数组 @ +2616,计数 @ +2624,步长 120 B
compile_options4CompileOptionsProto(消息)0x20CompileOptions::ToProtoCompileOptions @ +8
target_arguments5TargetArgumentsProto(消息)0x40内联构建(拓扑 + 配置属性),见下文
host_executions8repeated HostExecutionProto0x02数组 @ +2640,计数 @ +2648,步长 80 B
source_uri9string(已验证 UTF-8)0x04TpuExecutable::GetSourceUri,复制到 proto +56

怪异点 — 字段号并不连续。序列化器按标签升序发出 1、2、3、4、5、8、9;已填充集合中没有字段 6 或 7,因此假定 schema 是稠密 1..N 的重新实现会给 host_executions(8)和 source_uri(9)打错标签。source_uri 的标签在 0xf8d0c00 第 167 行的内联字符串写入器中以原始字节 740x4A = 字段 9,wire type 2)出现。

target_arguments(字段 5)的子装配

target_arguments 不是从单个源字段复制而来;它在 SerializeExecutable 内部就地构造为一个 TargetArgumentsProto,并由三个来源填充:

c
// SerializeExecutable @ 0xf8a9300, lines 422-568
target = proto.mutable_target_arguments();           // has_bit 0x40

// (a) topology, distilled then ToProto'd
TpuTopologySerdes::Distill(&serdes, exe->topology);  // topology @ +2536
serdes.ToProto(target->mutable_topology_args());     // TpuTopologyArgsProto, sub has_bit 0x04

// (b) configured properties
exe->configured_properties.ToProto(                  // @ +? via TpuConfiguredProperties
    target->mutable_configured_properties());        // sub has_bit 0x01

// (c) a scalar copied verbatim
target->field_48 = *(int*)(exe + 2576);              // sub has_bit 0x08

如果存在 MultiSliceConfigexe + 1832 非空),该路径还会反序列化一个 AotMegaScaleMultiSliceConfig,并把 MultiSliceTopologyAndLocationProto 写入 target_arguments(子 has_bit 0x02)。当跳过该分支时,会直接走 CompileOptions::ToProto(第 569-644 行)。拓扑/多切片路径上的失败会通过 AddSourceLocationImpl 转换成 absl::Status,位置在 tpu_pjrt_executable.cc 的第 421/439/445/448/453 行;这些行号锚定了每个子步骤。


AOT 容器 — xla::DeepseaExecutableProto

目的

DeepseaExecutableProto 是编译器内部容器,包装机器程序,即逐核 LLO/ISA bundle 和编译器自己的簿记信息。它是信封的字段 1,也是 AOT 编译器和 C-API 序列化/反序列化往返使用的独立单元。DeepseaExecutable::ToProto (0x134282e0) 和 FromProto (0x134283e0) 是它的转换器。

字段映射

来自 DeepseaExecutable::ToProto (0x134282e0):

字段标签类型has_bit来源
core_program1tpu::TpuCoreProgramProto(消息)0x01TpuCoreProgram::ToProtoTpuCoreProgram @ DeepseaExecutable+96
compiler_metadata2xdb::CompilerMetadata(消息)0x02Message::CopyFrom,来源 @ DeepseaExecutable+104

说明 — xdb::CompilerMetadataplatforms_deepsea::jellyfish::xdb::CompilerMetadata。它就是在整个 TpuJitResult/TpuCompilationCacheEntry::Program 生命周期中与程序并行携带的同一个元数据对象(构造函数位于 0xf7bbd600xf7bc5000xf8b7720),保存重新加载的程序所需、但不在 ISA 自身中的数据;buffer assignment、alias config 和 host-compute descriptors 在这些构造函数中作为它的兄弟对象被跟踪。

SerializeExecutable 构建信封字段 1 时,它并不调用 DeepseaExecutable::ToProto;而是把相同的两个子字段直接内联写入信封的 mutable_deepsea_executable()(第 144-377 行):把 TpuCoreProgram::ToProto 写入子字段 1,并对 exe+2584 处的 metadata 做一次 CopyFrom 写入子字段 2。独立的 ToProto 和内联路径会产生相同的双字段消息。


逐核 Bundle — tpu::TpuCoreProgramProto

目的

TpuCoreProgramProto 是单个 TPU 核心的已编译程序的序列化形式,即 TPU-to-LLO lowering 的输出。它是三个核心 dialect 上的 tagged union,外加一小块共享元数据。它的序列化器是 tpu::TpuCoreProgramProto::_InternalSerialize (0x1e82ea40)。

字段映射

标签和存在性从 0x1e82ea40 处的序列化器读取。has_bits 位于对象偏移 +16

字段标签(字节)Wire typehas_bit含义
(字段 2)16 = 0x10varint0x10标量选择器 / 计数(4 字节字段 @ +64 低字)
(字段 3)26 = 0x1ALEN0x02字节/字符串 blob(原始程序镜像)
(字段 4)32 = 0x20varint0x08varint(this+56
TensorCore/BarnaCore/SparseCore5 / 6 / 7LEN(消息)oneof @ +88core-program oneof;case = field_88 ∈ {5,6,7}
(字段 9)int64varint0x40int64(this+72)— 见 WriteInt64ToArrayWithField<9>
(字段 10)80 = 0x50varint(bool)0x20bool(this+68,仅当值为 1 时发出)
(字段 9,msg)byte_9[3]LEN(消息)0x04尾部子消息(this+48

怪异点 — oneof 由 v20 = *((_DWORD*)this + 22) 以及 0x1e82ea40 第 134-137 行的测试 (unsigned)(v20 - 5) <= 2 分派,也就是说,活动 arm 是判别量持有的 5/6/7 标签之一,并且只会写出 TpuCoreProgramProto_TensorCore (0x1e82dae0)、_BarnaCore (0x1e82dde0)、_SparseCore (0x1e82e240) 中的一个。重新实现必须把 5/6/7 视为互斥,而不是三个独立的可选消息。

子循环中的 repeated message(标签 byte_8this+24)和字段 3 的原始 blob 一起携带实际指令流;周围的 varint 携带核心类型判别量和大小。这个 blob 在 PJRT 反序列化期间不会被重新解析,它是加载时交给 TpuChip::NewProgramLoaded (0xe72b320) 的不透明 LLO/ISA 镜像。


分帧 — 长度分隔的片段

目的

TpuExecutableProtoToString (0xf8a8880) 是实际产出线缆字节的位置。它不会对整个 TpuExecutableProto 调用 proto.SerializeToString()。相反,它会写入一个 StringOutputStream,形式是有序的一串分隔(带长度前缀)子消息,把最大的 payload 从父消息中剥离出来,并最后序列化被耗减后的父消息。

算法

c
function TpuExecutableProtoToString(out, proto):          // 0xf8a8880
    stream = StringOutputStream(&result)

    // Segment 1: the per-core LLO bundle
    SerializeDelimitedToZeroCopyStream(                    // length-prefixed
        proto.deepsea_executable.core_program, stream)
    proto.deepsea_executable.clear_core_program()         // peel it out

    // Segment 2: compiler metadata
    SerializeDelimitedToZeroCopyStream(
        proto.deepsea_executable.compiler_metadata, stream)
    proto.deepsea_executable.clear_compiler_metadata()

    // Invariant: the inner container must now be empty
    if proto.deepsea_executable.ByteSizeLong() != 0:
        return Error("Not all fields in DeepseaExectuable "
                     "have been serialized")               // line 272

    // Segment 3: the HLO module
    SerializeDelimitedToZeroCopyStream(
        proto.hlo_module_with_config, stream)
    proto.clear_hlo_module_with_config()

    // Segment 4: the reduced parent (host transfers/executions,
    //            compile_options, target_arguments, source_uri)
    SerializeDelimitedToZeroCopyStream(proto, stream)      // "Reduced Proto"

    out = OK(move(result))

每个 SerializeDelimitedToZeroCopyStream 都会写出一个 varint 字节长度前缀,随后跟着消息字节,因此输出是自描述的:读取器按顺序拉取四个分隔帧,即 core program、compiler metadata、HLO module、reduced envelope。每个失败点都会在错误字符串中命名出问题片段的字节大小("Core program byte size is:"、"Compiler Metadata byte size is:"、"HLO Module byte size is"、"Reduced Proto byte size is:"),这是二进制中识别这种分帧的诊断指纹。

陷阱 — 因为父消息在写出最终帧前会被修改clear_core_programclear_compiler_metadataclear_hlo_module_with_config),线缆上的“reduced proto”缺少字段 1 的子项以及整个字段 2。天真的反序列化器如果对拼接结果做一次 ParseFromString,只会读到第一个分隔帧,并静默丢掉其余内容。读取器必须消费所有四个帧,并用片段 1 和 2 重新拼合内部 DeepseaExecutableProto

说明 — 这个拆分是规避 2 GiB 的措施。Protobuf 把单个序列化消息限制在 INT_MAX 字节;对于大型融合模型,LLO bundle 或 HLO proto 本身就可能接近该上限。把每个部分作为独立分隔帧发出,可以让每次 SerializeToArray 调用都保持在上限之下,同时仍然产出一个连续字节串。错误字符串("is it too large?")确认了大小上限正是动机。

反序列化路径

存在三个读取器,位于三个层级,并且它们与四片段写入器并不对称

text
PjRtClient::DeserializeExecutable (0xe6ecfa0) ── base class: returns
    "Deserializing serialized executable not supported." (unimplemented)

DeepseaCompiler::DeserializeExecutable (0xfaa2660) ── AOT path:
    DeepseaExecutableProto p;
    p.ParseFromString(bytes)                  // inner container only
    DeepseaExecutable::FromProto(p)           // 0x134283e0
    on failure: "Failed to parse serialized executable."
                (deepsea_compiler.cc:165)

TpuExecutable_Deserialize (0xeabede0) ── C runtime shim:
    DeepseaExecutableProto p;
    p.ParseFromArray(bytes, len)              // inner container only
    DeepseaExecutable::FromProto(p)
    on failure: "TpuExecutable_Deserialize: proto deserialization failed"

陷阱 — AOT/C-API 反序列化器解析的是 DeepseaExecutableProto不是 TpuExecutableProtoToString 产生的完整分隔式 TpuExecutableProto 流。这是两种不同的序列化形式,只共享内部容器:AOT 路径(DeepseaCompiler)只把编译器容器作为普通 protobuf 往返,而 PJRT 路径(TpuExecutableProtoToString)把它连同 HLO module、compile options 和 target block 一起包装进四片段分隔信封。重新实现者必须为产生这些字节的写入器选择匹配的读取器;它们不能互换。


C-API 程序 Blob — GetTpuProgramResponseExternal_Blob

目的

运行时 C-API 会通过比上述两个容器都更窄的包装器,在缓存/RPC 边界上传送单个核心的程序。TpuProgram_SerializeTpuExecutable (0xe8be720) 接收一个 TpuCoreProgram,将其序列化,并把结果包装进 tensorflow::tpu::GetTpuProgramResponseExternal_Blob

算法

c
function TpuProgram_SerializeTpuExecutable(core_program, out_bytes, out_status): // 0xe8be720
    blob = GetTpuProgramResponseExternal_Blob()
    s = blob.mutable_blob()                                 // ArenaStringPtr
    if !core_program.SerializeToString(s):                  // MessageLite
        out_status = Error("Failed to serialize proto, "
                           "invalid executable buffer.")    // tpu_program_c_api.cc:201
        return
    size = blob.ByteSizeLong()
    buf  = operator new(size)
    CHECK(blob.SerializePartialToArray(buf, size))          // proto_helper.h:45
    out_bytes = {buf, size}                                 // caller owns buf

配套读取器 TpuProgram_DeserializeFromGetTpuProgramResponseProto (0xe8be960) 和通用的 stream_executor::tpu::DeserializeProto<T, TpuSerializedProto> 模板(例如为 barna_core::HbmBuffersConfig0xf767580 处实例化)闭合了这一层的循环。

说明 — 尽管名称如此,TpuProgram_SerializeTpuExecutable 序列化的是 TpuCoreProgram,不是 TpuExecutable。它是逐核线缆单元;完整可执行文件的序列化是 TpuExecutable::SerializeExecutable (0xf8a9300)。配套的 TpuProgram_SerializeCompilerMetadata (0xe8be840) 对 CompilerMetadata 半边做同样的包装。


版本和兼容性戳

什么把序列化程序固定到目标上

序列化可执行文件只能在与它编译目标芯片代际匹配的硬件上加载。这个固定关系由两种机制携带:

  1. 目标块(target_arguments,字段 5)。SerializeExecutable(第 432-568 行)中的装配方式,它嵌入一个蒸馏后的 TpuTopologyArgsProto(来自 TpuTopologySerdes::Distill)和一个 TpuConfiguredPropertiesProtoTpuConfiguredProperties::ToProto)。二者共同记录程序构建时所针对的拓扑和已解析硬件属性。重新加载时,这里的不匹配会让来自某个 TPU 代际的程序无法在另一个代际上使用。

  2. 编译时版本 flag。 静态初始化器 _GLOBAL__sub_I_tpu_version_flag.cc (0x21367f50) 从 learning/45eac/tpu/runtime/tpu_version_flag.cc 注册三个 Abseil 命令行 flag:

Flag 符号作用
FLAGS_deepsea_versionTPU 代际 / 版本选择器
FLAGS_deepsea_variant某个代际内的硬件变体
FLAGS_deepsea_chip_config_name已命名的芯片配置

这些 flag 选择编译器构建所针对的 Target,解析后的值会进入 TpuConfiguredPropertiestarget_arguments。它们是该戳的上游来源,而不是线缆格式中的独立字段。

说明 —TpuExecutableProto 自身内部没有发现独立的 magic number、格式版本整数或 build-id 字段。兼容性是通过结构来执行的,即嵌入的拓扑/属性块以及 protobuf 自己的生成代码版本握手。重新实现应把 target_arguments 块视为兼容性戳;没有需要检查的独立单整数版本字段。(已在所调查的序列化/反序列化函数中确认缺失;由于并非每个嵌套 proto 都被穷尽追踪,因此“schema 中任何地方都不存在这种字段”的置信度为 MEDIUM。)

陷阱 — 因为戳是拓扑/属性块而不是廉价整数,所以快速“这个 blob 能在这里加载吗?”检查不能通过读取一个头部 word 完成。读取器必须把 reduced-envelope 帧解析到足以到达 target_arguments。廉价的序列化前兼容性门禁是缓存键,而不是序列化 blob;见 compilation-cache.md


对象布局参考

序列化器从活动 xla::TpuExecutableSerializeExecutable 中的实参 a2)读取的源偏移。偏移相对于对象基址;已对照 0xf8a9300 确认。

偏移字段用途
+8CompileOptions信封字段 4(compile_options
+1374flag byte控制是否以默认值发出 TpuCoreProgram/CompilerMetadata
+1504flag bytesource-URI / core-program 分支选择器
+1832MultiSliceConfig*字段 5 中的多切片拓扑子装配
+2536TpuTopology*(shared_ptr target)target_arguments.topology_args
+2560HloModule*信封字段 2(hlo_module_with_config
+2576int32target_arguments 标量(子字段 8)
+2584CompilerMetadata内部 DeepseaExecutableProto.compiler_metadata
+2600TpuCoreProgram*内部 DeepseaExecutableProto.core_program
+2616 / +2624HostTransferProto[] / count信封字段 3,步长 120 B
+2640 / +2648HostExecutionProto[] / count信封字段 8,步长 80 B

序列化基础设施函数

函数地址作用
xla::TpuExecutable::SerializeExecutable0xf8a9300从活动可执行文件构建信封
xla::TpuExecutableProtoToString0xf8a8880四片段分隔式分帧
xla::TpuExecutableProto::_InternalSerialize0xf8d0c00信封 proto 线缆写入器
xla::DeepseaExecutable::ToProto0x134282e0AOT 容器写入器
xla::DeepseaExecutable::FromProto0x134283e0AOT 容器读取器
tpu::TpuCoreProgramProto::_InternalSerialize0x1e82ea40逐核 bundle 写入器
tpu::TpuCoreProgramProto_TensorCore::_InternalSerialize0x1e82dae0oneof arm 5
tpu::TpuCoreProgramProto_BarnaCore::_InternalSerialize0x1e82dde0oneof arm 6
tpu::TpuCoreProgramProto_SparseCore::_InternalSerialize0x1e82e240oneof arm 7
xla::jellyfish::DeepseaCompiler::DeserializeExecutable0xfaa2660AOT 反序列化入口
xla::PjRtClient::DeserializeExecutable0xe6ecfa0基类(未实现)
TpuExecutable_Serialize0xeabea80C-API:DeepseaExecutable → proto
TpuExecutable_Deserialize0xeabede0C-API:bytes → DeepseaExecutable
TpuProgram_SerializeTpuExecutable0xe8be720C-API:core program → response blob
TpuProgram_SerializeCompilerMetadata0xe8be840C-API:metadata → blob
TpuProgram_DeserializeFromGetTpuProgramResponseProto0xe8be960C-API:response proto → program
tensorflow::tpu::internal::TpuCompileOpKernelImpl::BuildExecutable0xe8c0be0compile → TpuProgramGroup 生产者

交叉引用

  • 编译缓存 — 负责缓存键、fingerprint 和查找路径;廉价的序列化前兼容性门禁在那里,而不在 blob 中。
  • 编译器概览 — 序列化在端到端编译流水线中的位置。
  • 编译阶段 — 产生本页所序列化 TpuCoreProgram 的阶段顺序。
  • TPU-to-LLO ODS — 发出绑定在 TpuCoreProgramProto 内部的逐核 ISA 程序的 lowering。
  • PJRT 可执行文件执行 — 消费反序列化程序的 reload-and-enqueue 路径(TpuExecutable_LoadProgramAndEnqueueToStream)。
  • PJRT 概览 — 此格式往返通过的 PJRT Executable::Serialize/DeserializeExecutable C-API 表面。