序列化与格式

本文引用的文件 - xla.proto - xla_data.proto - protobuf_util.h - protobuf_util.cc - hlo.proto - hlo_dumps.md - debug_options_flags.h - debug_options_flags.cc - autotuner_cache.proto - autotuner.cc - file_based_autotuner_cache.cc - hlo_benchmark_runner.cc - metadata_collector.cc - metadata_utils.h - core/host_offloading/host_offloading_executable.proto - pjrt_c_api_shardings_extension.h - pjrt_layout_serdes.proto - executable_metadata.proto - xla_host_callback.proto - xla_sharding.proto - stream_executor_executable.proto - autotune_results.proto - autotuning.proto - autotune_result_wrapper.cc - autotune_result_wrapper.h - text_literal_reader.h - text_literal_writer.h - text_literal_reader.cc - text_literal_writer.cc - packed_literal_reader.h - packed_literal_reader.cc - sort_json.h - sort_json.cc - shape_util.proto - shape_util.h - shape_util.cc - shape_util_test.cc - shape_pool.h - shape_pool.cc - shape_pool_test.cc - literal_util.h - literal_util.cc - literal_util_test.cc - literal_pool.h - literal_pool.cc - literal_pool_test.cc - runtime/large_hlo_snapshot_serialization/large_hlo_snapshot_serialization.h - runtime/large_hlo_snapshot_serialization/large_hlo_snapshot_serialization.cc

目录

1. 引言
2. 项目结构
3. 核心组件
4. 架构总览
5. 详细组件分析
6. 依赖关系分析
7. 性能考量
8. 故障排查指南
9. 结论
10. 附录

引言

本文件面向XLA中间表示（HLO）的序列化与格式系统，系统性阐述以下内容： - HLO计算图的序列化格式与二进制编码规范 - 可移植API设计与跨平台传输能力 - HLO模块的元数据存储、版本管理与兼容性策略 - 序列化过程中的完整性校验、压缩与性能优化 - 反序列化的错误处理、回退机制与调试支持 - 不同格式之间的转换规则与限制条件 - 序列化与反序列化的最佳实践与性能建议

项目结构

围绕序列化与格式的关键目录与文件如下： - 协议缓冲区定义：xla.proto、xla_data.proto、service/hlo.proto、各后端与工具的proto - 序列化工具：protobuf_util.、large_hlo_snapshot_serialization. - 文本/打包读写：text_literal_reader/writer.、packed_literal_reader/writer. - 元数据与调试：hlo_dumps.md、debug_options_flags.、metadata_collector.、metadata_utils.h - 自动调优缓存：autotuner_cache.proto、autotuner.cc、file_based_autotuner_cache.cc、autotune_result_wrapper.*

graph TB
subgraph "协议定义"
XLA["xla.proto"]
DATA["xla_data.proto"]
HLO["service/hlo.proto"]
AUTOCACHE["backends/autotuner/autotuner_cache.proto"]
AUTORES["autotune_results.proto"]
AUTOCONF["autotuning.proto"]
end
subgraph "序列化工具"
PBUTIL_H["protobuf_util.h"]
PBUTIL_CC["protobuf_util.cc"]
LARGE_SNAP["runtime/large_hlo_snapshot_serialization.*"]
end
subgraph "文本与打包"
TXTREAD["text_literal_reader.*"]
TXTWR["text_literal_writer.*"]
PACKR["packed_literal_reader.*"]
end
subgraph "元数据与调试"
DUMPS["docs/hlo_dumps.md"]
DBGFLAGS["debug_options_flags.*"]
METACOL["backends/profiler/cpu/metadata_collector.cc"]
METAUTIL["backends/profiler/cpu/metadata_utils.h"]
end
XLA --- DATA
XLA --- HLO
AUTOCACHE --- AUTORES
AUTOCACHE --- AUTOCONF
PBUTIL_H --> PBUTIL_CC
LARGE_SNAP --- HLO
TXTREAD --- DATA
TXTWR --- DATA
PACKR --- DATA
DUMPS --- HLO
DBGFLAGS --- XLA
METACOL --- HLO
METAUTIL --- HLO

图表来源 - xla.proto - xla_data.proto - hlo.proto - autotuner_cache.proto - autotune_results.proto - autotuning.proto - protobuf_util.h - protobuf_util.cc - large_hlo_snapshot_serialization.h - large_hlo_snapshot_serialization.cc - text_literal_reader.h - text_literal_writer.h - packed_literal_reader.h - hlo_dumps.md - debug_options_flags.h - debug_options_flags.cc - metadata_collector.cc - metadata_utils.h

章节来源 - xla.proto - xla_data.proto - hlo.proto - hlo_dumps.md

核心组件

• 协议缓冲区定义
• xla.proto：包含调试选项、执行选项、设备句柄等，支撑跨平台编译与运行配置
• xla_data.proto：定义基础类型、形状、布局、字面量、收集通信参数等核心数据结构
• service/hlo.proto：定义HLO模块、指令、配置等服务层结构
• 序列化工具
• protobuf_util.*：提供基于序列化的消息比较与哈希，确保一致性与稳定性
• large_hlo_snapshot_serialization.*：针对大型快照的序列化封装
• 文本与打包读写
• text_literal_reader/writer.*：文本形式的字面量读写
• packed_literal_reader.*：紧凑二进制字面量读取
• 元数据与调试
• hlo_dumps.md：描述HloProto等转储格式
• debug_options_flags.*：调试标志与兼容性控制
• metadata_collector.*、metadata_utils.h：将HloProto注入到分析/可视化流程

章节来源 - xla.proto - xla_data.proto - protobuf_util.h - protobuf_util.cc - hlo_dumps.md - debug_options_flags.h - debug_options_flags.cc - metadata_collector.cc - metadata_utils.h

架构总览

下图展示XLA序列化与格式系统在编译与运行阶段的交互关系。

graph TB
A["前端/框架"] --> B["HLO模块构建<br/>service/hlo.proto"]
B --> C["序列化为HloProto<br/>xla_data.proto/LiteralProto"]
C --> D["持久化/传输<br/>文件/网络"]
D --> E["反序列化/加载<br/>protobuf_util.*"]
E --> F["执行/分析<br/>DebugOptions/ExecutionOptions"]
F --> G["元数据采集与可视化<br/>metadata_collector.cc"]
F --> H["自动调优缓存<br/>autotuner_cache.proto"]

图表来源 - hlo.proto - xla_data.proto - protobuf_util.h - protobuf_util.cc - metadata_collector.cc - autotuner_cache.proto

详细组件分析

组件A：HLO模块与指令序列化（HloProto）

• 结构要点
• HLO模块通过service/hlo.proto定义，包含指令列表、入口点、配置等
• 字面量通过xla_data.proto中的LiteralProto进行序列化，覆盖多种原生类型与稀疏编码
• 编码规范
• 使用Protocol Buffers标准二进制编码；字段顺序与默认值遵循proto3语义
• 字面量按类型分字段存储，小端序编码，稀疏索引与值分离
• 跨平台传输
• 通过xla.proto中的ExecutionOptions携带DebugOptions，确保跨平台行为一致
• 支持压缩转储（如xla_dump_compress_protos），降低带宽占用

sequenceDiagram
participant FE as "前端"
participant HLO as "HLO模块(service/hlo.proto)"
participant SER as "序列化(xla_data.proto)"
participant IO as "文件/网络"
participant DES as "反序列化(protobuf_util.*)"
FE->>HLO : 构建HLO模块
HLO->>SER : 生成HloProto/LiteralProto
SER->>IO : 写入二进制/文本
IO-->>DES : 读取序列化数据
DES-->>HLO : 反序列化为内存对象

图表来源 - hlo.proto - xla_data.proto - protobuf_util.h - protobuf_util.cc

章节来源 - hlo.proto - xla_data.proto - xla.proto

组件B：字面量序列化与文本/打包读写

• 字面量序列化
• LiteralProto覆盖布尔、整数、浮点、复数、稀疏索引等字段
• 小端序存储，FP16/BF16/U16/S16等字段明确编码方式
• 文本读写
• text_literal_reader/writer.*提供人类可读文本格式，便于调试与审计
• 打包读取
• packed_literal_reader.*用于高性能场景下的紧凑二进制读取

flowchart TD
Start(["开始"]) --> CheckFmt["选择读写格式<br/>文本/打包/二进制"]
CheckFmt --> Text["文本读写<br/>text_literal_reader/writer"]
CheckFmt --> Packed["打包读取<br/>packed_literal_reader"]
CheckFmt --> Binary["二进制读写<br/>LiteralProto"]
Text --> End(["结束"])
Packed --> End
Binary --> End

图表来源 - xla_data.proto - text_literal_reader.h - text_literal_writer.h - packed_literal_reader.h

章节来源 - xla_data.proto - text_literal_reader.cc - text_literal_writer.cc - packed_literal_reader.cc

组件C：元数据与调试（HloProto转储与可视化）

• 转储格式
• hlo_dumps.md指出HloProto为结构化转储格式，支持文本/二进制/HTML等多种输出
• 元数据采集
• metadata_collector.cc将HloProto注入到分析平面，metadata_utils.h提供统计与可视化接口
• 调试选项
• debug_options_flags.*提供大量调试开关，影响序列化/转储行为与验证策略

sequenceDiagram
participant RUN as "运行时"
participant DUMP as "转储(hlo_dumps.md)"
participant META as "元数据(metadata_collector.cc)"
participant VIS as "可视化(metadata_utils.h)"
RUN->>DUMP : 生成HloProto转储
DUMP->>META : 注入HloProto
META->>VIS : 提供统计/可视化

图表来源 - hlo_dumps.md - metadata_collector.cc - metadata_utils.h - debug_options_flags.h

章节来源 - hlo_dumps.md - metadata_collector.cc - metadata_utils.h - debug_options_flags.cc

组件D：自动调优结果的序列化与缓存

• 结构
• autotuner_cache.proto定义缓存键/值序列化结构
• autotuner.cc在自动调优过程中序列化结果并写入缓存
• file_based_autotuner_cache.cc提供文件级序列化/反序列化实现
• autotune_result_wrapper.*对键/值进行确定性序列化
• 版本与兼容性
• 通过proto字段编号与保留字段维持向后兼容
• 支持不同后端与平台的序列化差异

sequenceDiagram
participant AT as "自动调优器(autotuner.cc)"
participant WRAP as "结果包装(autotune_result_wrapper.*)"
participant CACHE as "缓存(file_based_autotuner_cache.cc)"
participant PROTO as "autotuner_cache.proto"
AT->>WRAP : 序列化键/值
WRAP-->>PROTO : 确定性序列化
PROTO->>CACHE : 写入/读取缓存
CACHE-->>AT : 返回序列化结果

图表来源 - autotuner.cc - autotuner.cc - autotuner_cache.proto - file_based_autotuner_cache.cc - autotune_result_wrapper.cc

章节来源 - autotuner.cc - autotuner.cc - autotuner_cache.proto - file_based_autotuner_cache.cc - autotune_result_wrapper.cc

组件E：大型快照序列化（大体积HLO）

• 需求背景
• 大型HLO模块在转储/传输时可能超出单文件或内存限制
• 实现思路
• large_hlo_snapshot_serialization.*提供分块/流式序列化能力，结合压缩与校验
• 与HloProto配合，支持超大规模图的稳定落盘与传输

flowchart TD
A["HLO模块"] --> B["分块/流式序列化"]
B --> C["压缩/校验"]
C --> D["持久化/传输"]
D --> E["恢复/反序列化"]

图表来源 - large_hlo_snapshot_serialization.h - large_hlo_snapshot_serialization.cc

章节来源 - large_hlo_snapshot_serialization.h - large_hlo_snapshot_serialization.cc

依赖关系分析

• 协议依赖
• xla.proto导入xla_data.proto与service/hlo.proto，形成“配置-数据-模块”的层级
• 各后端与工具（如PJRT、自动调优）通过proto导入使用核心数据结构
• 工具依赖
• protobuf_util.*被广泛用于一致性校验与哈希
• 文本/打包读写与字面量序列化紧密耦合
• 运行期依赖
• DebugOptions/ExecutionOptions贯穿序列化/转储/执行链路
• 元数据采集与可视化依赖HloProto注入

graph LR
XLA["xla.proto"] --> DATA["xla_data.proto"]
XLA --> HLO["service/hlo.proto"]
DATA --> TXT["text_literal_reader/writer.*"]
DATA --> PACK["packed_literal_reader.*"]
DATA --> PBUTIL["protobuf_util.*"]
HLO --> DUMP["hlo_dumps.md"]
HLO --> META["metadata_collector.cc"]
HLO --> AUTOCACHE["autotuner_cache.proto"]

图表来源 - xla.proto - xla_data.proto - hlo.proto - protobuf_util.h - text_literal_reader.h - text_literal_writer.h - packed_literal_reader.h - hlo_dumps.md - metadata_collector.cc - autotuner_cache.proto

章节来源 - xla.proto - xla_data.proto - hlo.proto - protobuf_util.h

性能考量

• 序列化开销
• 使用确定性序列化（protobuf_util.*）避免字段顺序差异导致的不必要比较
• 对大型HLO采用分块/流式序列化与压缩，降低I/O与内存峰值
• 压缩策略
• 支持转储压缩（如xla_dump_compress_protos），显著减少磁盘占用
• 读写路径
• 文本格式适合调试与审计；打包/二进制格式适合高吞吐传输
• 并发与缓存
• 自动调优缓存序列化结果，避免重复计算；文件级缓存提供快速命中

[本节为通用性能讨论，无需特定文件引用]

故障排查指南

• 反序列化错误
• 检查序列化格式一致性（protobuf_util.*提供比较与哈希工具）
• 核对字段编号与保留字段，避免破坏向后兼容
• 转储问题
• 确认hlo_dumps.md中格式选择与目标环境匹配
• 使用debug_options_flags.*调整验证级别与日志级别
• 自动调优缓存
• 若缓存损坏，清理后重新生成；确认autotuner_cache.proto字段编号正确
• 使用autotune_result_wrapper.*确保键/值序列化确定性

章节来源 - protobuf_util.h - protobuf_util.cc - hlo_dumps.md - debug_options_flags.h - autotuner_cache.proto - autotune_result_wrapper.cc

结论

XLA的序列化与格式系统以Protocol Buffers为核心，结合确定性序列化、压缩与分块策略，在保证跨平台一致性的同时兼顾性能与可维护性。通过完善的元数据采集与调试工具链，系统实现了从编译到执行的全链路可观测与可控。

[本节为总结性内容，无需特定文件引用]

附录

• 最佳实践
• 优先使用二进制序列化进行生产传输，文本格式仅用于调试
• 对大型HLO启用分块/流式序列化与压缩
• 使用protobuf_util.*进行一致性校验，避免因字段顺序差异引发的问题
• 在跨平台部署中，严格管理xla.proto中的调试与执行选项，确保行为一致
• 转换规则与限制
• 字面量序列化严格遵循xla_data.proto的字段与编码约定
• HloProto与文本/打包格式之间存在严格的解析与生成规则
• 自动调优缓存的键/值序列化需保持确定性，避免运行时差异

[本节为通用指导，无需特定文件引用]