このガイドは、ハードウェアをターゲットとするプログラムを XLA に出力して効率的に出力する必要があるシステム エンジニアを対象としています。詳細な手順ではなく、LLVM、Bazel、XLA に関する知識があることを前提としています。
XLA は、XLA によって出力された ML プログラムを実行するためのバックエンドを作成するために、新しいアーキテクチャまたはアクセラレータを実装できる抽象インターフェースを提供します。新しいハードウェア用に PyTorch や TensorFlow などのフロントエンド フレームワークから既存の op を実装するよりも、XLA のリターゲティングが大幅にシンプルかつスケーラブルになります。
ほとんどの実装は、次のいずれかのシナリオに分類されます。
- XLA でまだ正式にサポートされていない既存の CPU アーキテクチャ（既存の LLVM バックエンドの有無にかかわらず）。
- 既存の LLVM バックエンドを備えた、CPU に似たハードウェア。
- 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 ドメインに固有のものです。他の種類のハードウェア、たとえば Hexagon（アップストリーム LLVM バックエンドを持つ）などの DSP は、LLVM IR 出力ロジックの一部を再利用できますが、他の部分は一意です。
シナリオ 3: 既存の LLVM バックエンドがない、CPU に似たハードウェア
LLVM を使用できない場合は、目的のハードウェアに XLA 用の新しいバックエンドを実装することをおすすめします。このオプションでは最も手間がかかります。実装する必要があるクラスは次のとおりです。
StreamExecutor: 多くのデバイスでは、StreamExecutorのすべてのメソッドが必要というわけではありません。詳細については、既存のStreamExecutorの実装をご覧ください。xla::Compiler: このクラスは、HLO 計算のコンパイルをxla::Executableにカプセル化します。xla::Executable: このクラスは、プラットフォームでコンパイルされた計算を開始するために使用されます。xla::TransferManager: このクラスを使用すると、バックエンドは、所定のデバイスのメモリハンドルから XLA リテラルデータを作成するためのプラットフォーム固有のメカニズムを提供できます。つまり、ホストとデバイスの間のデータ転送をカプセル化するのに役立ちます。