Questa pagina è stata tradotta dall'API Cloud Translation.
Switch to English

Sviluppo di un nuovo backend per XLA

Questa guida preliminare è destinata ai primi utenti che desiderano facilmente effettuare il retargeting di TensorFlow sul proprio hardware in modo efficiente. La guida non è passo per passo e presuppone la conoscenza di LLVM , Bazel e TensorFlow.

XLA fornisce un'interfaccia astratta che una nuova architettura o acceleratore può implementare per creare un back-end per eseguire grafici TensorFlow. Il retargeting XLA dovrebbe essere significativamente più semplice e scalabile rispetto all'implementazione di ogni TensorFlow Op esistente per il nuovo hardware.

La maggior parte delle implementazioni rientrerà in uno dei seguenti scenari:

  1. Architettura della CPU esistente non ancora ufficialmente supportata da XLA, con o senza un back-end LLVM esistente.
  2. Hardware non simile alla CPU con un back-end LLVM esistente.
  3. Hardware non simile alla CPU senza un back-end LLVM esistente.

Scenario 1: architettura CPU esistente non ancora ufficialmente supportata da XLA

In questo scenario, iniziare guardando il backend della CPU XLA esistente. XLA semplifica il retarget di TensorFlow su CPU diverse utilizzando LLVM, poiché la differenza principale tra i backend XLA per CPU è il codice generato da LLVM. Google testa XLA per architetture x64 e ARM64.

Se il fornitore dell'hardware ha un back-end LLVM per il proprio hardware, è semplice collegare il back-end con LLVM creato con XLA. In modalità JIT, il backend della CPU XLA emette il codice per la CPU host. Per la compilazione xla::AotCompilationOptions , xla::AotCompilationOptions può fornire una tripla LLVM per configurare l'architettura di destinazione.

Se non esiste un back-end LLVM esistente ma esiste un altro tipo di generatore di codice, dovrebbe essere possibile riutilizzare la maggior parte del back-end CPU esistente.

Scenario 2: hardware non simile alla CPU con un back-end LLVM esistente

È possibile modellare una nuova implementazione di xla::Compiler sulle xla::CPUCompiler esistenti xla::CPUCompiler e xla::GPUCompiler , poiché già emettono LLVM IR. A seconda della natura dell'hardware, è possibile che molti aspetti della generazione IR di LLVM debbano essere modificati, ma un sacco di codice può essere condiviso con i backend esistenti.

Un buon esempio da seguire è il backend GPU di XLA. Il backend GPU si rivolge a un ISA non simile alla CPU e pertanto alcuni aspetti della sua generazione di codice sono unici per il dominio GPU. Altri tipi di hardware, ad esempio DSP come Hexagon (che ha un back-end LLVM a monte), possono riutilizzare parti della logica di emissione IR LLVM, ma altre parti saranno uniche.

Scenario 3: hardware non simile alla CPU senza un back-end LLVM esistente

Se non è possibile utilizzare LLVM, l'opzione migliore è implementare un nuovo back-end per XLA per l'hardware desiderato. Questa opzione richiede il massimo sforzo. Le classi che devono essere implementate sono le seguenti:

  • StreamExecutor : per molti dispositivi non sono necessari tutti i metodi di StreamExecutor . Vedi le implementazioni esistenti di StreamExecutor per i dettagli.
  • xla::Compiler : questa classe incapsula la compilazione di un calcolo xla::Executable in un xla::Executable .
  • xla::Executable : questa classe viene utilizzata per avviare un calcolo compilato sulla piattaforma.
  • xla::TransferManager : questa classe consente ai backend di fornire meccanismi specifici della piattaforma per la costruzione di dati letterali XLA da determinati handle di memoria del dispositivo. In altre parole, aiuta a incapsulare il trasferimento di dati dall'host al dispositivo e viceversa.