此页面由 Cloud Translation API 翻译。

使用 XLA 工具

XLA 开发工作流通常围绕 HLO IR，表示隔离功能提供给编译器的计算。XLA 附带多个命令行工具（如下所述），它们会使用 HLO 并运行 HLO，或提供中间编译阶段。使用此类工具 compile->modify->run 迭代周期，因为 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 运行，并且默认情况下将编译、运行和与参考解释器的比较捆绑在一起实施。例如，用于运行输入文件的常规调用 computation.hlo 并检查其正确性：

$ run_hlo_module --platform=CUDA --reference_platform=Interpreter computation.hlo

与所有工具一样，--help 可用于获取完整的选项列表。

运行支持 SPMD 的 HLO 代码段：`multihost_hlo_runner`

多主机 HLO 运行程序 (Multihost HLO runner) 是一款非常相似的工具，但要注意，它支持 SPMD，包括跨主机通信。请参阅如需了解详情，请参阅多主机 HLO 运行程序。

多 HLO 重放

run_hlo_module 和 hlo_runner_main，这通常有助于重放转储中的所有模块目录：

$ hlo_runner_main /dump/*before_optimizations*

正在运行 HLO 编译的通过/阶段：`hlo-opt`

在调试或了解编译器的运行情况时，获取特定硬件在流水线（无论是 HLO、经过优化的 HLO、TritonIR 还是 LLVM），针对给定（稳定）HLO 输入。

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

该命令会将转储输出到 stdout（如果指定了 -o，则输出到给定文件）。

无设备使用情况

大部分编译都不需要访问 GPU，并且可以指定我们可以在命令行上获取 GPU 规范，例如：无权访问加速器：

$ 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 --xla-hlo-enable-passes-only=algebraic_simplifer input.hlo