测试框架¶

目录¶

简介
项目结构
核心组件
架构总览
详细组件分析
依赖关系分析
性能考量
故障排查指南
结论
附录

简介¶

本文件系统性梳理 XLA 的测试体系，覆盖单元测试、集成测试、HLO 传递测试与多主机 HLO 运行器（Multi-Host HLO Runner）的架构、测试用例编写规范与执行流程，并提供可操作的测试示例路径、测试数据准备、断言方法与结果验证技巧。同时记录持续集成与性能回归检测相关实践，帮助开发者高效编写与维护高质量测试。

项目结构¶

XLA 的测试分布在多个层次： - 文档层：提供 HLO 传递测试与多主机 HLO 运行器使用指南 - 单元测试层：以 GoogleTest/GMock 为主，覆盖 HLO 传递流水线、后端特定变换等 - 集成测试层：通过多主机 HLO 运行器在真实硬件或仿真环境下执行复杂 HLO 模块 - 工具与数据层：提供运行器、测试数据集与构建脚本

graph TB
subgraph "文档"
D1["docs/test_hlo_passes.md"]
D2["docs/tools_multihost_hlo_runner.md"]
end
subgraph "单元测试"
U1["xla/hlo/pass/hlo_pass_pipeline_test.cc"]
U2["xla/backends/gpu/transforms/alias_passthrough_params_test.cc"]
U3["xla/tests/xla_test_backend_predicates.{h,cc}"]
end
subgraph "集成测试工具"
T1["xla/tools/multihost_hlo_runner/functional_hlo_runner.cc"]
T2["xla/tools/multihost_hlo_runner/create_client.cc"]
T3["xla/tools/multihost_hlo_runner/BUILD"]
end
subgraph "测试数据"
DATA["xla/tools/multihost_hlo_runner/data/*.hlo<br/>*.pbtxt"]
end
D1 --> U1
D1 --> U2
D2 --> T1
D2 --> T2
T1 --> DATA
T2 --> DATA
U3 --> U1
U3 --> U2

图示来源 - docs/test_hlo_passes.md - docs/tools_multihost_hlo_runner.md - xla/hlo/pass/hlo_pass_pipeline_test.cc - xla/backends/gpu/transforms/alias_passthrough_params_test.cc - xla/tests/xla_test_backend_predicates.h - xla/tests/xla_test_backend_predicates.cc - xla/tools/multihost_hlo_runner/functional_hlo_runner.cc - xla/tools/multihost_hlo_runner/create_client.cc - xla/tools/multihost_hlo_runner/BUILD

章节来源 - docs/test_hlo_passes.md - docs/tools_multihost_hlo_runner.md

核心组件¶

HLO 传递测试规范与工具链：基于 FileCheck/LIT 的 HLO 传递测试方法，支持 CHECK 前缀与自动化生成脚本
多主机 HLO 运行器：支持单进程/多进程、单节点/多节点复现 HLO，便于集成测试与回归验证
后端谓词与参数化测试：通过环境变量驱动的后端谓词，实现跨设备/后端的条件跳过与参数化
HLO 传递流水线单元测试：验证模块级传递的变更、不变与不变式检查器行为

章节来源 - docs/test_hlo_passes.md - docs/tools_multihost_hlo_runner.md - xla/tests/xla_test_backend_predicates.h - xla/tests/xla_test_backend_predicates.cc - xla/hlo/pass/hlo_pass_pipeline_test.cc

架构总览¶

下图展示了测试框架的整体架构：文档指导层、测试实现层、工具与数据层之间的关系，以及关键执行路径。

graph TB
A["HLO 传递测试规范<br/>docs/test_hlo_passes.md"] --> B["GoogleTest/GMock 测试套件"]
B --> C["HLO 传递流水线单元测试<br/>hlo_pass_pipeline_test.cc"]
B --> D["后端特定变换测试<br/>alias_passthrough_params_test.cc"]
E["多主机 HLO 运行器<br/>tools/multihost_hlo_runner"] --> F["运行器入口<br/>functional_hlo_runner.cc"]
E --> G["客户端创建<br/>create_client.cc"]
E --> H["测试数据集<br/>data/*.hlo, *.pbtxt"]
I["后端谓词库<br/>xla_test_backend_predicates.{h,cc}"] --> B
I --> E

图示来源 - docs/test_hlo_passes.md - xla/hlo/pass/hlo_pass_pipeline_test.cc - xla/backends/gpu/transforms/alias_passthrough_params_test.cc - xla/tools/multihost_hlo_runner/functional_hlo_runner.cc - xla/tools/multihost_hlo_runner/create_client.cc - xla/tools/multihost_hlo_runner/data/sharded_2_devices.hlo - xla/tests/xla_test_backend_predicates.h - xla/tests/xla_test_backend_predicates.cc

详细组件分析¶

HLO 传递测试（FileCheck/LIT）¶

测试风格
使用 FileCheck 在输入 HLO 中内联 CHECK 行，统一使用“// CHECK”作为分隔符
推荐将 CHECK 与输入 IR 局部对齐，提升可读性与可维护性
对于 GPU 测试，可结合 LIT 运行器与 hlo-opt，按后端目标生成对应前缀的 CHECK
自动化辅助
提供 generate_hlo_test_checks.py 脚本，自动插入 CHECK 指令，需人工复核一致性
不推荐的做法
避免遍历结果图进行逐节点匹配，此类测试难以阅读与调试

flowchart TD
Start(["开始：编写 HLO 输入"]) --> InlineCheck["内联 CHECK 指令<br/>使用 // CHECK"]
InlineCheck --> RunOpt["运行 hlo-opt 或 Pass 流水线"]
RunOpt --> Compare["比对输出与 CHECK"]
Compare --> AutoGen{"是否使用自动生成脚本？"}
AutoGen --> |是| ManualReview["人工复核生成的 CHECK"]
AutoGen --> |否| Proceed["直接使用手写 CHECK"]
ManualReview --> Pass{"全部匹配？"}
Proceed --> Pass
Pass --> |是| Ok["测试通过"]
Pass --> |否| Fix["修正 HLO/Pass 或更新 CHECK"]
Fix --> InlineCheck

图示来源 - docs/test_hlo_passes.md

章节来源 - docs/test_hlo_passes.md

多主机 HLO 运行器（集成测试）¶

功能概述
支持在单 GPU 或多 GPU 上运行 HLO 模块；可仅编译不运行
支持单进程与多进程（MPI/SLURM），单节点与多节点场景
关键参数
--run=false：仅编译，必要时可利用持久化自动调优缓存
--hlo_argument_mode=uninitialized：减少初始化开销与主机内存占用
--run_xla_backend_only 或 --xla_disable_all_hlo_passes：避免重复编译导致的错误
典型流程
从模型导出 HLO（before/after 优化）
使用运行器在目标拓扑上复现，校验数值与性能指标

sequenceDiagram
participant Dev as "开发者"
participant Dump as "HLO 导出"
participant Runner as "多主机 HLO 运行器"
participant Client as "客户端/分布式初始化"
participant GPU as "GPU 设备"
Dev->>Dump : "导出 HLObefore/after 优化"
Dump-->>Dev : "得到 .hlo / .pbtxt"
Dev->>Runner : "传入 HLO 文件与参数"
Runner->>Client : "初始化 PJRT/NCCL 客户端"
Client->>GPU : "分配设备/拓扑"
Runner->>GPU : "编译并执行 HLO"
GPU-->>Runner : "返回结果/时间"
Runner-->>Dev : "输出日志/指标"

图示来源 - docs/tools_multihost_hlo_runner.md - xla/tools/multihost_hlo_runner/functional_hlo_runner.cc - xla/tools/multihost_hlo_runner/create_client.cc

章节来源 - docs/tools_multihost_hlo_runner.md

后端谓词与参数化测试¶

目标
基于环境变量（XLA_TEST_DEVICE、XLA_TEST_DEVICE_TYPE、XLA_TEST_MODIFIERS）判断当前后端，动态启用/禁用测试或参数化输入
常用谓词
DeviceIs/DeviceTypeIs/HasModifiers/BackendLike/BackendIsExactly/BackendIsStrict
BackendSupportsFloat64/BackendSupportsComplex128
Empty() 用于参数化测试中生成空输入集合
使用建议
在测试套件中通过谓词包裹条件逻辑，必要时配合 GTEST_SKIP()/GTEST_ALLOW_UNINSTANTIATED_PARAMETERIZED_TEST()

flowchart TD
Env["读取环境变量<br/>XLA_TEST_*"] --> Pred["谓词函数计算"]
Pred --> Decision{"满足条件？"}
Decision --> |是| Enable["启用测试/参数化"]
Decision --> |否| Skip["跳过测试/空参数化"]

图示来源 - xla/tests/xla_test_backend_predicates.h - xla/tests/xla_test_backend_predicates.cc

章节来源 - xla/tests/xla_test_backend_predicates.h - xla/tests/xla_test_backend_predicates.cc

HLO 传递流水线单元测试¶

测试要点
验证模块级传递是否改变图结构（changed/unchanged）
异步线程执行流中的传递应用范围
不变式检查器（Invariant Checker）在流水线中的失败传播
传递元数据（Pass Metadata）的设置与校验
示例参考
模块名重命名传递、根指令名称反转、混合流水线、不变式检查器失败、无变更场景

sequenceDiagram
participant Test as "测试用例"
participant Pipeline as "HloPassPipeline"
participant Pass1 as "FooToBarModulePass"
participant Pass2 as "BazToQuxModulePass"
participant Checker as "BarBlowerUpper(不变式)"
Test->>Pipeline : "AddPass/FooToBarModulePass"
Test->>Pipeline : "AddPass/BazToQuxModulePass"
Test->>Pipeline : "AddInvariantChecker/BarBlowerUpper"
Pipeline->>Pass1 : "RunImpl(module)"
Pass1-->>Pipeline : "changed=true/false"
Pipeline->>Pass2 : "RunImpl(module)"
Pass2-->>Pipeline : "changed=true/false"
Pipeline->>Checker : "RunImpl(module)"
Checker-->>Pipeline : "OK 或返回错误"
Pipeline-->>Test : "返回状态与变更标志"

图示来源 - xla/hlo/pass/hlo_pass_pipeline_test.cc

章节来源 - xla/hlo/pass/hlo_pass_pipeline_test.cc

后端特定变换测试（示例：GPU 别名透传参数）¶

目标
验证别名透传参数变换在 tuple 输出中的行为，确保不会对同一参数建立多重别名
断言要点
通过 input_output_alias_config 获取别名映射，断言参数到参数的别名关系
预设别名场景下的行为一致性

flowchart TD
Parse["解析 HLO 模块"] --> Apply["运行 AliasPassthroughParams 变换"]
Apply --> Config["读取 input_output_alias_config"]
Config --> Assert1["断言 p0 -> p0 映射"]
Config --> Assert2["断言输出未别名 p1若适用"]
Assert1 --> Done["测试完成"]
Assert2 --> Done

图示来源 - xla/backends/gpu/transforms/alias_passthrough_params_test.cc

章节来源 - xla/backends/gpu/transforms/alias_passthrough_params_test.cc

依赖关系分析¶

组件耦合
HLO 传递测试依赖文档规范与工具链（FileCheck/LIT/generate_hlo_test_checks.py）
多主机运行器依赖 PJRT/NCCL 客户端与测试数据集
后端谓词为所有测试提供统一的后端感知能力
外部依赖
GoogleTest/GMock 用于单元测试
Bazel 构建系统管理测试目标与运行器构建

graph LR
Docs["HLO 传递测试规范"] --> Tests["单元测试"]
Tools["多主机运行器"] --> IntTests["集成测试"]
Pred["后端谓词库"] --> Tests
Pred --> IntTests
Tests --> Data["测试数据集"]
IntTests --> Data

图示来源 - docs/test_hlo_passes.md - docs/tools_multihost_hlo_runner.md - xla/tests/xla_test_backend_predicates.h - xla/tools/multihost_hlo_runner/BUILD

章节来源 - docs/test_hlo_passes.md - docs/tools_multihost_hlo_runner.md - xla/tools/multihost_hlo_runner/BUILD