XLA ツールの使用

XLA 開発ワークフローは通常、コンパイラに渡される分離された関数計算を表す HLO IR を中心に行われます。XLA には、HLO を使用して実行するか、中間コンパイル ステージを提供する複数のコマンドライン ツール(後述)が付属しています。HLO は可視化とハッキングが可能なため、このようなツールを使用すると、compile->modify->run の反復処理サイクルを迅速に行うことができます。多くの場合、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_modulehlo_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

これにより、ダンプが 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 のフラグもサポートされているため、このツールを使用して 1 回のパスの実行をテストできます。

$ hlo-opt --platform=CUDA --stage=hlo --passes=algebraic_simplifer input.hlo