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

本予備ガイドは、効率的な方法でTensorFlowをハードウェアに容易に対応させたいと考えている、アーリーアダプターのためのものです。 本ガイドは1つ1つ丁寧に説明したものではなく、LLVM、Bazel、TensorFlowの知識を前提としています。

XLAは、新しいアーキテクチャやアクセラレータが、TensorFlowのグラフを処理するバックエンドを実装するための抽象的なインターフェースを提供します。 XLAへの対応は、新しいハードウェア向けに既存のすべてのTensorFlowのオペレーションを実装するのと比べて、はるかに簡潔でスケーラブルです。

実装のほとんどは、以下のシナリオのうちの1つに分類されます。

  1. LLVMのバックエンドが存在するかしないかにかかわらず、公式にXLAでサポートされていない既存のCPUアーキテクチャ
  2. LLVMのバックエンドが存在する、CPUではないハードウェア
  3. LLVMのバックエンドが存在しない、CPUではないハードウェア

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

このシナリオの場合、既存の XLA CPUバックエンド を見ることから始めてください。 XLAのCPUバックエンド間の主な違いは、LLVMによって生成されるコードであることから、XLAではLLVMを使って異なるCPUをTensorFlowに簡単に対応できます。 Googleは、x64とARM64のアーキテクチャに対してXLAを試験しています。

もしハードウェア企業がハードウェア向けのLLVMのバックエンドをもつ場合、ビルドされたLLVMのバックエンドをXLAに接続することは簡単です。 JITモードでは、XLAのCPUバックエンドはホスト側のCPUのコードを生成します。 Ahead-Of-Timeコンパイルでは、xla::AotCompilationOptions が対象とするアーキテクチャに対して設定するLLVM Tripleを提供します。

もし既存のLLVMのバックエンドがなくてもコード生成器が違う形で存在するならば、既存のCPUバックエンドの大部分を再利用できる可能性があります。

シナリオ2: LLVMのバックエンドが存在する、CPUではないハードウェア

LLVM IRを出力する既存の xla::CPUCompiler や xla::GPUCompiler クラスをベースとして、新しい xla::Compiler の実装を作ることが可能です。 ハードウェアの性質によりLLVM IRの生成方法は異なりますが、多くのコードは既存のバックエンドと共有できるでしょう。

よい参考例は、XLAの GPUバックエンド です。 GPUのバックエンドはCPUとは異なるISAをターゲットとするため、GPUドメイン固有なコードの生成方法になります。 ほかの種類のハードウェア、たとえば（アップストリームのLLVMのバックエンドをもつ）HexagonのようなDSPは、LLVM IRの生成のしくみを再利用することができますが、ほかの部分は固有のものになるでしょう

シナリオ3: LLVMのバックエンドが存在しない、CPUではないハードウェア

LLVMを利用できない場合、対象のハードウェア向けに新しいバックエンドを実装することが最良の選択肢となります。 この選択肢は、多大な労力を必要とします。 実装しなければならないクラスは次に示すとおりです。

  • StreamExecutor: 多くのデバイスでは、StreamExecutor のすべてのメソッドが必要になることはありません。詳細は既存の StreamExecutor の実装を見てください。
  • xla::Compiler: 本クラスは、HLO Computationから xla::Executable へのコンパイル処理を隠蔽します。
  • xla::Executable: 本クラスは、コンパイル済みのComputationをプラットフォーム上で実行するために使用されます。
  • xla::TransferManager: 本クラスは、与えられたデバイスメモリのハンドルからXLAのリテラルデータを構築するための、プラットフォーム特有のしくみを提供することを可能にします。言い換えれば、ホストからデバイスまたはその反対のデータ転送処理を隠蔽します。