Sviluppo di un nuovo backend per XLA

Questa guida è rivolta agli ingegneri di sistema che desiderano che XLA produca programmi che hanno come target il loro hardware in modo efficiente. La guida non contiene istruzioni dettagliate e presuppone la conoscenza di LLVM, Bazel e XLA.

XLA offre un'interfaccia astratta che può essere implementata da una nuova architettura o un nuovo acceleratore per creare un backend per l'esecuzione di programmi ML come output da XLA. Il retargeting XLA dovrebbe essere significativamente più semplice e scalabile rispetto all'implementazione di ogni operazione 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 della 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 XLA esistente. XLA semplifica il targeting di CPU diverse utilizzando LLVM, perché 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 un LLVM triplo 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 di LLVM IR debbano essere modificati, ma una grande quantità di codice può essere condivisa con i backend esistenti.

Un buon esempio da seguire è il backend GPU di XLA. Il backend 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 DSP come Hexagon (che ha un backend LLVM upstream), possono riutilizzare parti della logica di emissione IR di LLVM, ma altre parti saranno univoche.

Scenario 3: hardware senza CPU senza un backend LLVM esistente

Se non è possibile utilizzare LLVM, l'opzione migliore è implementare un nuovo backend per XLA per l'hardware desiderato. Questa opzione richiede il massimo impegno. 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 specificati. In altre parole, aiuta a incapsulare il trasferimento dei dati dall'host al dispositivo e viceversa.