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 Runner」。
多 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 檔案是 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