この記事は、The Parallel Universe Magazine 58 号に掲載されている「JAX and OpenXLA Part 2: Run Process and Underlying Logic」の日本語参考訳です。原文は更新される可能性があります。原文と翻訳文の内容が異なる場合は原文を優先してください。
パート 1 では、JAX が Python 表現を StableHLO 表現に変換する方法や、StableHLO ダンプから HLO (高水準演算) への変換など、JAX フレームワーク (英語) と OpenXLA* 向けインテル® エクステンション (英語) の基本概念について説明しました。パート 2 は実行ワークフローの続きで、HLO ダンプから LLVM IR (中間表現) を作成し、最終的にインテルの GPU で実行可能な SPIR-V ファイルを生成します。
実際の実行ロジック
HLO から LLVM IR への変換
図 1. PJRT プラグインが実装された OpenXLA* 向けインテル® エクステンション（出典 (英語))
HLO 表現から LLVM 中間表現 (IR) への変換において、インテルの OpenXLA コンパイラーは、高水準演算表現から低水準マシンコードを生成します。HLO はテンソル演算を基本単位として使用し、モデルに含まれる数学演算とテンソル操作を表現します。詳細は、LLVM ドキュメント (英語) を参照してください。
HLO は、OpenXLA コンパイラーで使用される高水準中間表現であり、ディープラーニング・ モデルの計算グラフを表します。LLVM IR は、マシンコード抽象化に近い低水準中間表現であり、通常は特定の命令レベルの演算の記述に使用されます。LLVM IR は LLVM コンパイラー・フレームワークの中核であり、さまざまなコンパイラー・ツールチェーンで広く使用されています。
HLO と LLVM IR
- 変換プロセス: OpenXLA のコンパイルフローにおいて、HLO は高水準の中間表現です。一連の最適化を経て、最終的に LLVM IR 表現に変換されます。この変換プロセスでは通常、HLO の高水準演算を改良し、それらを低水準の命令セットにマッピングします。
- マッピングと改良:
- 演算マッピング: 各 HLO 演算は、複数の LLVM IR 命令に対応する場合があります。例えば、行列乗算の HLO 演算は、LLVM IR の一連のロード命令、乗算命令、累算命令に展開される場合があります。
- ハードウェア固有の最適化: LLVM IR の生成中に、コンパイラーは特定のハードウェア機能 (ベクトル化、パイプライン化など) を活用して、さらなる最適化を行い、効率良いコードを生成する場合があります。
- コード生成: HLO が LLVM IR に変換されると、LLVM IR は LLVM バックエンド・ツールチェーンを介して処理され、最終的にハードウェア・プラットフォーム固有のマシンコード (x86、Arm、CUDA コードなど) が生成されます。ただし、インテルの GPU では最初に SPIR-V に変換されます。このマシンコードはターゲットデバイス上で実行されます。
- HLO と LLVM IR の関係性: HLO と LLVM IR の関係は、高水準テンソル計算表現から低水準マシンコード生成への橋渡しと見ることができます。HLO は、プラットフォームに依存しない高水準抽象化を提供し、グローバル最適化を容易にします。一方、LLVM IR はこれらの抽象化を特定の命令セットに変換し、ハードウェア・レベルの最適化と実行を可能にします。OpenXLA コンパイラーはこのプロセスにより、ディープラーニング・モデルをさまざまなハードウェア・プラットフォームに効率的にマッピングし、最適な計算性能を実現します。
ダンプ LLVM ファイルの環境変数
すべてのパスのログをダンプするには、「–xla_dump_hlo_pass_re=.*」 オプションを追加する必要があります。
$export XLA_FLAGS="--xla_dump_hlo_as_text --xla_dump_to=./dump"
module_0001.jit_integer_pow.ir-no-opt.ll (初期 LLVM IR）
module_0000.jit_convert_element_type.ir-with-opt.ll (最終 LLVM IR）
module_0001.jit_integer_pow.spv (SPIR-V ファイル）
module_0001.jit_integer_pow.thunk_sequence.txt (実行順序）
図 2. HLO ダンプから LLVM への変換に関連するファイルの一覧
