Développement d'un nouveau backend pour XLA

Ce guide préliminaire est destiné aux premiers utilisateurs qui souhaitent recibler facilement TensorFlow sur leur matériel de manière efficace. Le guide n'est pas détaillé et suppose une connaissance de LLVM , Bazel et TensorFlow.

XLA fournit une interface abstraite qu'une nouvelle architecture ou un nouvel accélérateur peut implémenter pour créer un backend pour exécuter des graphes TensorFlow. Le retargeting XLA devrait être beaucoup plus simple et évolutif que la mise en œuvre de chaque opération TensorFlow existante pour un nouveau matériel.

La plupart des implémentations tomberont dans l'un des scénarios suivants :

1. Architecture CPU existante non encore officiellement prise en charge par XLA, avec ou sans backend LLVM existant.
2. Matériel non semblable à un processeur avec un backend LLVM existant.
3. Matériel non semblable à un processeur sans backend LLVM existant.

Scénario 1 : architecture de processeur existante non encore officiellement prise en charge par XLA

Dans ce scénario, commencez par examiner le backend CPU XLA existant. XLA permet de recibler facilement TensorFlow sur 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. Google teste XLA pour les architectures x64 et ARM64.

Si le fournisseur de matériel dispose d'un backend LLVM pour son matériel, il est simple de lier le backend au LLVM construit avec XLA. En mode JIT, le backend CPU XLA émet du code pour le CPU hôte. Pour une compilation anticipée, xla::AotCompilationOptions peut fournir un triplet LLVM pour configurer l'architecture cible.

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

Scénario 2 : matériel non similaire à un processeur avec un backend LLVM existant

Il est possible de modéliser une nouvelle xla::Compiler sur les classes existantes xla::CPUCompiler et xla::GPUCompiler , puisque celles-ci émettent déjà LLVM IR. Selon la nature du matériel, il est possible que de nombreux aspects de la génération IR LLVM devront ê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 GPU cible un ISA non similaire au CPU, et donc certains aspects de sa génération de code sont uniques au domaine GPU. D'autres types de matériel, par exemple des DSP comme Hexagon (qui a un backend LLVM en amont), peuvent réutiliser des parties de la logique d'émission IR LLVM, mais d'autres parties seront uniques.

Scénario 3 : matériel non similaire à un processeur sans backend LLVM existant

S'il n'est pas possible d'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 de StreamExecutor ne sont pas nécessaires. Voir les implémentations existantes de StreamExecutor pour plus de détails.
• xla::Compiler : Cette classe encapsule la compilation d'un calcul HLO dans un xla::Executable .
• xla::Executable : Cette classe permet de lancer un calcul compilé sur la plateforme.
• xla::TransferManager : cette classe permet aux backends de fournir des mécanismes spécifiques à la plate-forme pour la construction de données littérales XLA à partir de poignées de mémoire de périphérique données. En d'autres termes, cela aide à encapsuler le transfert de données de l'hôte vers l'appareil et inversement.