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 ファイルが指定されたフォルダに保存されます。また、他の多くの有用なアーティファクトも保存されます。

[run_hlo_module] HLO モジュールを実行する

bazel run //xla/tools:run_hlo_module -- [flags] <filename>

ツール run_hlo_module は最適化前の HLO で動作し、デフォルトではコンパイル、実行、参照インタープリタ実装との比較をバンドルします。たとえば、NVIDIA GPU で入力ファイル computation.hlo を実行し、その正しさを確認する通常の呼び出しは次のようになります。

run_hlo_module --platform=CUDA --reference_platform=Interpreter computation.hlo

複数の HLO モジュールを実行する

run_hlo_module では、複数の HLO モジュールを使用した呼び出しがサポートされています。ディレクトリ内のすべての hlo モジュールを実行するには:

bazel run //xla/tools:run_hlo_module -- [flags] /dump/*before_optimizations*

[multihost_hlo_runner] SPMD サポートで HLO モジュールを実行

# Note: Binary name is `hlo_runner_main`.
bazel run //xla/tools/multihost_hlo_runner:hlo_runner_main -- [flags] <filename>

マルチホスト HLO ランナーは、クロスホスト通信を含む SPMD をサポートするという点で、非常に類似したツールです。詳しくは、マルチホスト HLO ランナーをご覧ください。

SPMD サポートで複数の HLO モジュールを実行する

run_hlo_module と同様に、multihost_hlo_runner も複数のモジュールでの呼び出しをサポートしています。

bazel run //xla/tools/multihost_hlo_runner:hlo_runner_main -- [flags] /dump/*before_optimizations*

[hlo-opt] HLO モジュールをコンパイル

bazel run //xla/tools:hlo-opt -- --platform=[gpu|cpu|...] [more flags] <filename>

コンパイラの動作をデバッグまたは理解する際に、特定の HLO または StableHLO 入力に対して、パイプラインの特定の時点（HLO、最適化された HLO、TritonIR、LLVM など）で特定のハードウェアの展開を取得すると便利なことがよくあります。

hlo-opt は、PTX、最適化後の HLO、最適化前の LLVM IR、TritonIR など、複数の出力ステージをサポートしています。サポートされているステージの正確なセットはプラットフォームによって異なります（たとえば、PTX は NVIDIA 固有です）。--list-stages コマンドを使用すると確認できます。

hlo-opt --platform=CUDA --list-stages
buffer-assignment
hlo
hlo-backend
html
llvm
llvm-after-optimizations
llvm-before-optimizations
ptx

ステージを選択すると、特定のプラットフォームの変換結果を特定のストリームに書き込むことができます。

hlo-opt --platform=cpu --stage=hlo input.hlo

これにより、ダンプが stdout に出力されます（-o が指定されている場合は、指定されたファイルに出力されます）。

GPU のデバイスレス コンパイル

デバイスレス コンパイルでは GPU にアクセスする必要はありません。Deviceless Compilation は、GPU へのアクセスが必要なステージでコマンドライン（--xla_gpu_target_config_filename）で GPU 仕様を指定する方法を提供し、GPU デバイスの必要性を排除します。

例: GPU デバイスにアクセスできない場合の PTX 出力:

hlo-opt  --platform=CUDA --stage=llvm  --xla_gpu_target_config_filename=/xla/tools/hlo_opt/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"

GPU の詳細な仕様については、/xla/tools/hlo_opt/gpu_specs をご覧ください。

自動チューニング

コンパイルには、コンパイル --stage に基づく自動チューニングが含まれることがあります。デバイスレス コンパイルを機能させるには、ユーザーが --xla_gpu_autotune_level=0
で自動チューニングを
無効にするか、--xla_gpu_load_autotune_results_from=<filename>（--xla_gpu_dump_autotune_results_to=<filename> で取得）で既存の自動チューニング結果を
読み込む必要があります。

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_to=<myfile.pbtxt> を使用してシリアル化できます。

[hlo-opt] HLO パスの開発とデバッグ

# If you are working with hardware independent passes from the
# `xla/hlo/transforms/` directory, prefer light-weight version
# of the `hlo-opt` tool with fewer dependencies:

bazel run //xla/hlo/tools:hlo-opt -- [flags] <filename>

# Otherwise, for hardware independent and CPU, GPU passes use
# the same binary from "Compile HLO Modules" section above:

bazel run //xla/tools:hlo-opt -- [flags] <filename>

hlo-opt ツールを使用すると、指定されたプラットフォームのコンパイル ステージとは独立して個々のパスを実行できます。この分離により、入力 hlo モジュールでパスをすばやく実行し、障害の根本原因を特定できます。

hlo-opt --passes=schedule-aware-collective-cse input.hlo

hlo-opt ツールは DebugOptions XLA_FLAGS もサポートしています。

hlo-opt --passes=schedule-aware-collective-cse
--xla_gpu_experimental_collective_cse_distance_threshold=20 input.hlo

--list-passes オプションを使用して、パス名文字列を取得します。

hlo-opt --list-passes

--passes オプションに複数のパスを指定することで、独自のカスタム パイプラインを作成できます。

hlo-opt --passes=pass1,pass2,pass3 input.hlo

新しい HLO パスの開発を支援する

  1. まず、パスを作成します。

  2. 新しいパスを hlo-opt ツールパス レジストリに登録します。

    RegisterPass<FooPass>(FooPassInputOptions)

    パスの種類に基づいて、登録する次のいずれかの場所を選択します。
    opt_lib.cc ハードウェアに依存しないパス。
    cpu_opt.cc CPU 固有のパス。
    gpu_opt.cc GPU 固有のパス。
    compiled_opt.cc CPU、GPU、XPU に共通のパス。
    ビルドの依存関係を追加することを忘れないでください。

    すべての hlo-opt ユーザーがパスを使用できるように、パスの登録を PR（例）の一部として含めます。

  3. hlo-opt ツールを再ビルドし、--list-passes オプションを使用してパスの登録が成功したことを検証してから、--passes オプションを使用してパスを実行します。

    $ hlo-opt --passes=foo-pass input.hlo

  4. パスの単体テストの作成については、https://openxla.org/xla/test_hlo_passes を参照してください。

パスのランタイム測定

大規模なモデルの場合、完全なコンパイルの実行に数分かかることがあり、微妙なパフォーマンスの低下を検出するのが困難になります。一方、hlo-opt を使用した個々のパス実行では、パフォーマンスを正確に測定し、新しいコード変更によって発生した実行時間のわずかな増加も簡単に検出できます。

time hlo-opt --passes=reduce-window-rewriter,scatter_simplifier
--xla_reduce_window_rewrite_base_length=128 input.hlo

[hlo-opt] HLO モジュール形式を変換

# Use the light weight version of the `hlo-opt` tool.

bazel run //xla/hlo/tools:hlo-opt -- [flags] <filename>

HLO Text -> HLO Proto に変換

hlo-opt --emit-proto input.hlo

HLO Proto または HLO Proto Binary -> HLO Text に変換

hlo-opt input.pbtxt or input.pb

[ptx-opt] コンパイラ LLVM モジュールから PTX まで

このツールは LLVMIR 最適化パイプラインを実行し、CompileToPtx を呼び出します。

bazel run //xla/hlo/tools/ptx-opt -- --arch=9.0 <filename>

このツールは、すべてのパスの後に LLVMIR をダンプすることもできます。

bazel run //xla/hlo/tools/ptx-opt -- --arch=9.0 --xla_dump_to=<path> --xla_gpu_dump_llvmir <filename>