XLA 用の新しいバックエンドの開発

このガイドは、ハードウェアを効率的にターゲットとするプログラムを XLA で出力したいシステム エンジニアを対象としています。このガイドでは、順を追って説明するのではなく、LLVM、Bazel、XLA に関する知識があることを前提としています。

XLA は、XLA によって出力された ML プログラムを実行するバックエンドを作成するために、新しいアーキテクチャまたはアクセラレータが実装できる抽象インターフェースを提供します。XLA のリターゲティングは、新しいハードウェア用に PyTorch や TensorFlow などのフロントエンド フレームワークから既存のすべてのオペレーションを実装するよりも、はるかにシンプルでスケーラブルである必要があります。

ほとんどの実装は、次のいずれかのシナリオに分類されます。

  1. 既存の LLVM バックエンドの有無にかかわらず、XLA でまだ正式にサポートされていない既存の CPU アーキテクチャ。
  2. 既存の LLVM バックエンドを備えた、CPU に似たハードウェア。
  3. 既存の LLVM バックエンドを持たない、CPU に似たハードウェア。

シナリオ 1: XLA でまだ正式にサポートされていない既存の CPU アーキテクチャ

このシナリオでは、まず既存の XLA CPU バックエンドを確認します。CPU の XLA バックエンドの主な違いは LLVM によって生成されるコードであるため、XLA では LLVM を使用してさまざまな CPU を簡単にターゲットに設定できます。

ハードウェア ベンダーがハードウェア用の LLVM バックエンドを使用している場合は、XLA で構築された LLVM にバックエンドを簡単にリンクできます。JIT モードでは、XLA CPU バックエンドはホスト CPU のコードを出力します。事前コンパイルの場合、xla::AotCompilationOptions で LLVM トリプルを提供して、ターゲット アーキテクチャを構成できます。

既存の LLVM バックエンドがなくても、別の種類のコード生成ツールがある場合は、既存の CPU バックエンドのほとんどを再利用できます。

シナリオ 2: 既存の LLVM バックエンドを使用した、CPU に似たハードウェア

既存の xla::CPUCompiler クラスと xla::GPUCompiler クラスは、すでに LLVM IR を送出しているため、新しい xla::Compiler 実装をモデル化できます。ハードウェアの性質によっては、LLVM IR 生成のさまざまな側面を変更する必要がありますが、多くのコードを既存のバックエンドと共有できます。

わかりやすい例は、XLA の GPU バックエンドです。GPU バックエンドは CPU に似た ISA をターゲットにしているため、コード生成の一部は GPU ドメインに固有のものです。他の種類のハードウェア（アップストリーム LLVM バックエンドを持つ Hexagon などの DSP など）では、LLVM IR 出力ロジックの一部を再利用できますが、他の部分は一意です。

シナリオ 3: 既存の LLVM バックエンドがない、CPU に似たハードウェア

LLVM を利用できない場合は、目的のハードウェアに XLA 用の新しいバックエンドを実装することをおすすめします。このオプションでは最も労力がかかります。実装する必要があるクラスは次のとおりです。

  • StreamExecutor: 多くのデバイスでは、StreamExecutor のすべてのメソッドが必要というわけではありません。詳細については、既存の StreamExecutor の実装をご覧ください。
  • xla::Compiler: このクラスは、HLO 計算のコンパイルを xla::Executable にカプセル化します。
  • xla::Executable: このクラスは、コンパイル済みの計算をプラットフォームで開始するために使用されます。
  • xla::TransferManager: このクラスを使用すると、バックエンドは特定のデバイスメモリ ハンドルから XLA リテラルデータを作成するためのプラットフォーム固有のメカニズムを提供できます。つまり、ホストとデバイス間のデータ転送をカプセル化するのに役立ちます。