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 ランナーをご覧ください。
マルチ HLO リプレイ
複数のモジュールでの呼び出しは、run_hlo_module
と
hlo_runner_main
: ダンプ内のすべてのモジュールを再生する場合に便利
ディレクトリ:
$ hlo_runner_main /dump/*before_optimizations*
HLO コンパイルのパス/ステージの実行: hlo-opt
コンパイラの動作をデバッグしたり理解したりするときには、 特定の時点における特定のハードウェアの拡張を 特定の(Stable)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_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