XLA (Accelerated Linear Algebra) è un compilatore di machine learning (ML) che ottimizza l'algebra lineare, migliorando la velocità di esecuzione e l'utilizzo della memoria. Questa pagina fornisce una breve panoramica degli obiettivi e dell'architettura del compilatore XLA.
Obiettivi
Attualmente, XLA supporta diversi frontend di framework ML (tra cui PyTorch, TensorFlow e JAX) e fa parte del progetto OpenXLA, un ecosistema di tecnologie di compilazione open source per ML sviluppate in modo collaborativo dalle principali organizzazioni software e hardware di ML. Prima della creazione del progetto OpenXLA, XLA è stato sviluppato all'interno del progetto TensorFlow, ma gli obiettivi fondamentali rimangono gli stessi:
Migliorare la velocità di esecuzione. Compila sottografici per ridurre i tempi di esecuzione delle operazioni di breve durata ed elimina l'overhead del runtime, fonde le operazioni con pipeline per ridurre l'overhead di memoria e specializza le forme di tensore note per consentire una propagazione costante più aggressiva.
Migliora l'utilizzo della memoria. Analizzare e pianificare l'utilizzo della memoria, eliminando molti buffer di archiviazione intermedi.
Ridurre la dipendenza da operazioni personalizzate. Elimina la necessità di molte operazioni personalizzate migliorando le prestazioni delle operazioni di basso livello fuse automaticamente per adeguarle a quelle delle operazioni personalizzate originariamente fuse a mano.
Migliorare la portabilità. Rendi relativamente facile scrivere un nuovo backend per hardware innovativo, in modo che una grande parte di modelli ML possa essere eseguita senza modifiche su quell'hardware. Questo è in contrasto con l'approccio specializzato sulle operazioni monolitiche individuali per il nuovo hardware, che richiede la riscrittura dei modelli per poterle utilizzare.
Come funziona
Il compilatore XLA prende i grafici dei modelli dai framework ML definiti in StableHLO e li compila nelle istruzioni delle macchine per varie architetture. StableHLO definisce un set di operazioni con controllo delle versioni (HLO = operazioni di alto livello) che fornisce un livello di portabilità tra i framework ML e il compilatore.
In generale, il processo di compilazione che converte il grafico del modello in un file eseguibile ottimizzato per il target prevede i seguenti passaggi:
XLA esegue sul grafico StabiliHLO diversi passaggi di ottimizzazione e analisi integrati che sono indipendenti dalla destinazione, ad esempio CSE, fusione operativa indipendente dal target e analisi del buffer per l'allocazione della memoria di runtime per il calcolo. Durante la fase di ottimizzazione, XLA converte anche il dialetto StableHLO in un dialetto HLO interno.
XLA invia il calcolo HLO a un backend per ulteriori ottimizzazioni a livello di HLO, questa volta tenendo presenti informazioni e esigenze specifiche per il target. Ad esempio, il backend GPU potrebbe eseguire fusioni delle operazioni vantaggiose in particolare per il modello di programmazione GPU e determinare come partizionare il calcolo in flussi. In questa fase, i backend possono anche associare pattern a determinate operazioni o combinazioni delle stesse alle chiamate alla libreria ottimizzate.
Il backend esegue quindi la generazione di codice specifico per la destinazione. I backend di CPU e GPU inclusi con XLA utilizzano LLVM per IR di basso livello, ottimizzazione e generazione di codice. Questi backend emettono l'IR di LLVM necessario per rappresentare il calcolo dell'HLO in modo efficiente e quindi richiamano il codice LLVM per emettere codice nativo da questo IR di LLVM.
In questo processo, il compilatore XLA è modulare, in quanto è facile inserire lo slot in un backend alternativo per scegliere come target alcune nuove architetture HW. Il backend GPU attualmente supporta le GPU NVIDIA tramite il backend LLVM NVPTX. Il backend della CPU supporta più ISA della CPU.