通常、XLA 開発ワークフローでは HLO IR が中心となります。HLO は、コンパイラに提供される独立した関数計算を表します。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_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