使用 XLA 工具

XLA 開發工作流程通常以 HLO IR 為中心,代表提供給編譯器的獨立功能運算。XLA 隨附多項指令列工具 (如下所述),可使用 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 運作,且根據預設,套件會編譯、執行,並與參考解譯器實作進行比較。舉例來說,一般叫用 NVIDIA GPU 執行輸入檔案 computation.hlo 並檢查其正確性:

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

和所有工具一樣,--help 也可用於取得完整的選項清單。

在支援 SPMD 的情況下執行 HLO 程式碼片段:multihost_hlo_runner

多主機 HLO 執行器是非常相似的工具,需要注意它支援 SPMD,包括跨主機通訊。一般叫用方式如下:

hlo_runner_main  --device_type=gpu --use_spmd_partitioning=true --num_partitions=4 --num_replicas=1 --hlo_file=computation.hlo

執行 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

它會將傾印列印至 stdout (如果指定 -o,則會列印至指定檔案)。

無裝置使用情形

大部分的編譯作業都不需要 GPU 存取權,但在指令列中指定 GPU 規格,我們可以取得這類來源,例如沒有加速器存取權的 PTX 輸出內容:

$ hlo-opt  --platform=CUDA --stage=llvm  --xla_gpu_target_config_filename=(pwd)/tools/data/gpu_specs/a100_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_80.txtpb --xla_gpu_load_autotune_results_from=results.textpb input.hlo

Autotune 檔案是 autotune_results.proto 的文字序列化,範例如下:

version: 2
results {
  device: "sm_8.0 with 42331013120B RAM, 108 cores, 1410000KHz clock, 1215000KHz mem clock, 41943040B L2$"
  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