使用 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 執行器

重播多何動作

run_hlo_module 和多個模組都支援叫用多個模組 hlo_runner_main,通常適合在傾印中重播所有模組 目錄:

$ hlo_runner_main /dump/*before_optimizations*

執行中通過/HLO 編譯的階段:hlo-opt

偵錯或瞭解編譯器的運作方式時,經常會很實用 即可讓特定硬體在 (例如 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_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