Ce guide est destiné aux ingénieurs système qui souhaitent que XLA génère des programmes qui ciblent efficacement son matériel. Ce guide n'est pas détaillé et part du principe 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 chaque opération existante à partir d'un framework d'interface tel que PyTorch ou TensorFlow pour le nouveau matériel.
La plupart des implémentations relèvent de l'un des scénarios suivants:
- Architecture de processeur existante pas encore officiellement compatible avec XLA, avec ou sans backend LLVM existant.
- Matériel ne ressemblant pas à un processeur avec un backend LLVM existant
- Matériel ne ressemblant pas à 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, il est simple d'associer le backend au LLVM créé avec XLA. En mode JIT, le backend de processeur XLA émet du code pour le processeur hôte. Pour une compilation anticipée, xla::AotCompilationOptions peut fournir un triple LLVM pour configurer l'architecture cible.
S'il n'existe pas de backend LLVM, mais qu'un autre type de générateur de code existe, il devrait être possible de réutiliser la majeure partie du backend de processeur existant.
Scénario 2: Matériel ne ressemblant 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 infrarouges LLVM. Selon la nature du matériel, il est possible que de nombreux aspects de la génération IR LLVM doivent être modifiés, mais une grande partie du code peut être partagée avec les backends existants.
Le backend GPU de XLA est un bon exemple à suivre. Le backend de GPU cible un ISA sans processeur. Par conséquent, certains aspects de sa génération de code sont propres au domaine 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 IR LLVM, mais d'autres sont uniques.
Scénario 3: Matériel ne ressemblant pas à un processeur, sans backend LLVM existant
Si vous ne pouvez pas utiliser LLVM, la meilleure option consiste à implémenter 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 d'StreamExecutorne sont pas nécessaires. Pour en savoir plus, consultez les implémentationsStreamExecutorexistantes.xla::Compiler: cette classe encapsule la compilation d'un calcul HLO dans un objetxla::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, elle permet d'encapsuler le transfert de données de l'hôte vers l'appareil, et inversement.