XLA 开发工作流通常以 HLO IR 为中心,后者表示提供给编译器的隔离函数计算。XLA 附带多个命令行工具(详见下文),这些工具会使用 HLO 并运行它,或提供中间编译阶段。使用此类工具对于快速完成 compile->modify->run
迭代周期至关重要,因为 HLO 既可直观呈现,又可进行黑客攻击,而迭代更改和运行 HLO 通常是了解和修正 XLA 性能或行为的最快方式。
如需获取使用 XLA 编译的程序的 HLO,最简单的方法通常是使用 XLA_FLAGS
环境变量:
$ XLA_FLAGS=--xla_dump_to=/tmp/myfolder ./myprogram-entry-point
该工具会将所有优化前 HLO 文件以及许多其他实用工件存储在指定的文件夹中。
运行 HLO 代码段:run_hlo_module
工具 run_hlo_module
对预优化 HLO 进行操作,默认情况下会将编译、运行和与参考解释器实现的比较打包在一起。例如,在 NVIDIA GPU 上运行输入文件 computation.hlo
并检查其正确性的常规调用如下所示:
$ run_hlo_module --platform=CUDA --reference_platform=Interpreter computation.hlo
与所有工具一样,--help
可用于获取完整的选项列表。
运行具有 SPMD 支持的 HLO 代码段:multihost_hlo_runner
多主机 HLO 运行程序是一个非常类似的工具,但需要注意的是,它支持 SPMD,包括跨主机通信。如需了解详情,请参阅多主机 HLO 运行程序。
多 HLO 重放
run_hlo_module
和 hlo_runner_main
都支持使用多个模块进行调用,这通常便于重放转储目录中的所有模块:
$ hlo_runner_main /dump/*before_optimizations*
正在运行 HLO 编译的各个阶段/步骤:hlo-opt
在调试或了解编译器的工作原理时,通常很有用的是,针对给定(稳定)HLO 输入,在流水线中的特定位置(无论是 HLO、经过优化的 HLO、TritonIR 还是 LLVM)获取特定硬件的展开。
hlo-opt
支持多个输出阶段:PTX、优化后的 HLO、优化前的 LLVM IR 或 TritonIR。支持的确切阶段集取决于平台(例如,PTX 是 NVIDIA 专有的),您可以使用 --list-stages 命令查看:
$ hlo-opt --platform=CUDA --list-stages
hlo
llvm
ptx
选择阶段后,用户可以将给定平台的转换结果写入给定数据流:
$ hlo-opt myinput.hlo --platform=CUDA --stage=llvm
该命令会将转储输出到标准输出(如果指定了 -o
,则会输出到给定文件)。
无设备使用
大多数编译不需要访问 GPU,通过在命令行中指定 GPU 规范,我们无需访问加速器即可获取 PTX 输出:
$ hlo-opt --platform=CUDA --stage=llvm --xla_gpu_target_config_filename=(pwd)/tools/data/gpu_specs/a100_pcie_80.txtpb input.hlo
编译器随附了常用 GPU 的规范,所提供的文件是 device_description.proto
的字符串序列化:
gpu_device_info {
cuda_compute_capability {
major: 8
minor: 0
}
threads_per_block_limit: 1024
threads_per_warp: 32
shared_memory_per_block: 127152
shared_memory_per_core: 65536
threads_per_core_limit: 2048
core_count: 6192
fpus_per_core: 64
block_dim_limit_x: 2147483647
block_dim_limit_y: 65535
block_dim_limit_z: 65535
memory_bandwidth: 2039000000000
l2_cache_size: 4194304
clock_rate_ghz: 1.1105
device_memory_size: 79050250240
}
platform_name: "CUDA"
如果需要自动调整,无设备编译可能会遇到问题。幸运的是,我们还可以在命令行中提供这些信息:
$ hlo-opt --platform=CUDA --stage=llvm --xla_gpu_target_config_filename=gpu_specs/a100_pcie_80.txtpb --xla_gpu_load_autotune_results_from=results.textpb input.hlo
自动调优文件是 autotune_results.proto
的文本序列化,示例如下:
version: 3
results {
device: "CUDA: 8.0, Cores: 108, GPU clock: 1.41 GHz, Memory bandwidth: 1555 GB/s, L2 cache: 40 MB"
hlo: "{\n tmp_0 = f16[1,16,17,3]{3,2,1,0} parameter(0)\n tmp_1 = f16[16,51]{1,0} bitcast(f16[1,16,17,3]{3,2,1,0} tmp_0)\n tmp_2 = s8[16,17,3]{2,1,0} parameter(1)\n tmp_3 = s8[51,16]{0,1} bitcast(s8[16,17,3]{2,1,0} tmp_2)\n tmp_4 = f16[51,16]{0,1} convert(s8[51,16]{0,1} tmp_3)\n tmp_5 = f16[16,16]{1,0} dot(f16[16,51]{1,0} tmp_1, f16[51,16]{0,1} tmp_4), lhs_contracting_dims={1}, rhs_contracting_dims={0}\n ROOT tmp_6 = f16[1,16,16]{2,1,0} bitcast(f16[16,16]{1,0} tmp_5)\n}"
result {
run_time {
nanos: 31744
}
triton {
block_m: 32
block_n: 32
block_k: 32
split_k: 1
num_stages: 1
num_warps: 4
}
}
}
可以使用 XLA_FLAGS=--xla_gpu_dump_autotune_results_t=<myfile.pbtxt>
对自动调整数据库进行序列化
运行单个编译器传递
该工具还支持 XLA_FLAGS
中的标志,因此可以用于测试运行单次传递:
$ hlo-opt --platform=CUDA --stage=hlo --passes=algebraic_simplifer input.hlo