Sviluppo di un nuovo backend per XLA

Questa guida è rivolta agli ingegneri di sistema che vogliono che XLA generi programmi che hanno come target il loro hardware in modo efficiente. La guida non è passo passo e presuppone la conoscenza di LLVM, Bazel e XLA.

XLA fornisce un'interfaccia astratta che può essere implementata da una nuova architettura o da un nuovo acceleratore per creare un backend per l'esecuzione di programmi ML tramite XLA. Il retargeting XLA dovrebbe essere molto più semplice e più scalabile rispetto all'implementazione di ogni operazione operativa esistente da un framework di frontend come PyTorch o TensorFlow per il nuovo hardware.

La maggior parte delle implementazioni rientra in uno dei seguenti scenari:

  1. Architettura CPU esistente non ancora ufficialmente supportata da XLA, con o senza un backend LLVM esistente.
  2. Hardware senza CPU con un backend LLVM esistente.
  3. Hardware senza CPU senza un backend LLVM esistente.

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

In questo scenario, inizia esaminando il backend CPU XL esistente. XLA semplifica il targeting di diverse CPU utilizzando LLVM, poiché la differenza principale tra i backend XLA per le CPU è il codice generato da LLVM.

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

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

Scenario 2: hardware senza CPU con un backend LLVM esistente

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

Un buon esempio da seguire è il backend GPU di XLA. Il backend della GPU ha come target un ISA senza CPU, pertanto alcuni aspetti della generazione del codice sono univoci per il dominio GPU. Altri tipi di hardware, ad esempio i DSP come Hexagon (che ha un backend LLVM a monte), possono riutilizzare parti della logica di emissione IR LLVM, ma altre parti sono univoche.

Scenario 3: hardware senza CPU senza un backend LLVM esistente

Se non è possibile utilizzare LLVM, la soluzione migliore è quella di implementare un nuovo backend 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. Per informazioni dettagliate, consulta le implementazioni esistenti di StreamExecutor.
  • xla::Compiler: questa classe racchiude la compilazione di un calcolo HLO 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 creare dati letterali XLA dagli handle della memoria del dispositivo. In altre parole, aiuta a incapsulare il trasferimento dei dati dall'host al dispositivo e viceversa.