Cette page a été traduite par l'API Cloud Translation.
Switch to English

Développer un nouveau backend pour XLA

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

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

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

  1. L'architecture CPU existante n'est pas encore officiellement prise en charge par XLA, avec ou sans backend LLVM existant.
  2. Matériel non-CPU avec un backend LLVM existant.
  3. Matériel non-CPU sans backend LLVM existant.

Scénario 1: l'architecture CPU existante n'est pas encore officiellement prise en charge par XLA

Dans ce scénario, commencez par examiner le backend du processeur XLA existant. XLA facilite le reciblage de TensorFlow vers 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 du processeur XLA émet du code pour le processeur hôte. Pour une compilation xla::AotCompilationOptions , xla::AotCompilationOptions peut fournir un triple 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-CPU avec un backend LLVM existant

Il est possible de modéliser une nouvelle xla::Compiler sur les xla::CPUCompiler existantes xla::CPUCompiler et xla::GPUCompiler , car celles-ci émettent déjà LLVM IR. En fonction de la nature du matériel, il est possible que de nombreux aspects de la génération LLVM IR devront être modifiés, mais beaucoup de code peut être partagé avec les backends existants.

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

Scénario 3: Matériel non-CPU 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 qui doivent être implémentées sont les suivantes:

  • StreamExecutor : Pour de nombreux appareils, toutes les méthodes de StreamExecutor sont pas nécessaires. Voir les StreamExecutor 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 est utilisée pour 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 construire des données littérales XLA à partir de descripteurs de mémoire de périphérique donnés. En d'autres termes, il aide à encapsuler le transfert des données de l'hôte vers l'appareil et inversement.