Développer un nouveau backend pour la XLA

Ce guide est destiné aux ingénieurs système qui souhaitent que XLA génère des programmes qui ciblent efficacement leur matériel. Il n'est pas détaillé et suppose que vous connaissez LLVM, Bazel et XLA.

XLA fournit une interface abstraite qu'une nouvelle architecture ou un nouvel accélérateur peut implémenter pour créer un backend permettant d'exécuter des programmes de ML générés par XLA. Le reciblage XLA devrait être beaucoup plus simple et plus évolutif que la mise en œuvre de toutes les opérations existantes à partir d'un framework d'interface tel que PyTorch ou TensorFlow pour le nouveau matériel.

La plupart des implémentations appartiennent à l'un des scénarios suivants:

  1. Architecture de processeur existante pas encore officiellement compatible avec XLA, avec ou sans backend LLVM existant.
  2. Matériel ne faisant pas appel à un processeur avec un backend LLVM existant
  3. Matériel ne faisant pas appel à un processeur sans backend LLVM existant

Scénario 1: Architecture de processeur existante pas encore officiellement compatible avec XLA

Dans ce scénario, commencez par examiner le backend de processeur XLA existant. XLA permet de cibler facilement différents processeurs à l'aide de LLVM, car la principale différence entre les backends XLA pour les processeurs est le code généré par LLVM.

Si le fournisseur de matériel dispose d'un backend LLVM pour son matériel, vous pouvez facilement l'associer au backend LLVM créé avec XLA. En mode JIT, le backend de processeur XLA émet du code pour le processeur hôte. Pour une compilation à l'avance, xla::AotCompilationOptions peut fournir un triple LLVM pour configurer l'architecture cible.

S'il n'existe pas de backend LLVM, mais qu'il existe un autre type de générateur de code, il devrait être possible de réutiliser la majeure partie du backend de processeur existant.

Scénario 2: Matériel qui ne ressemble pas à un processeur avec un backend LLVM existant

Il est possible de modéliser une nouvelle implémentation xla::Compiler sur les classes xla::CPUCompiler et xla::GPUCompiler existantes, car celles-ci émettent déjà des images LLVM IR. Selon la nature du matériel, il est possible que de nombreux aspects de la génération de flux infrarouges LLVM doivent être modifiés, mais une grande partie du code peut être partagée avec les backends existants.

Un bon exemple à suivre est le backend GPU de XLA. Le backend de GPU cible un ISA de type processeur. Par conséquent, certains aspects de sa génération de code sont uniques au domaine du GPU. D'autres types de matériel, tels que les DSP comme Hexagon (qui dispose d'un backend LLVM en amont), peuvent réutiliser des parties de la logique d'émission infrarouge LLVM, mais d'autres parties seront uniques.

Scénario 3: Matériel ne faisant pas appel à un processeur sans backend LLVM existant

S'il n'est pas possible d'utiliser LLVM, la meilleure option consiste à mettre en œuvre un nouveau backend pour XLA pour le matériel souhaité. Cette option demande le plus d'efforts. Les classes à implémenter sont les suivantes:

  • StreamExecutor : pour de nombreux appareils, toutes les méthodes de StreamExecutor ne sont pas nécessaires. Pour en savoir plus, consultez les implémentations StreamExecutor existantes.
  • xla::Compiler : cette classe encapsule la compilation d'un calcul HLO dans un élément xla::Executable.
  • xla::Executable : cette classe permet de lancer un calcul compilé sur la plate-forme.
  • xla::TransferManager : cette classe permet aux backends de fournir des mécanismes spécifiques à la plate-forme pour construire des données littérales XLA à partir de identifiants de mémoire d'appareil donnés. En d'autres termes, il permet d'encapsuler le transfert de données de l'hôte vers l'appareil, et inversement.