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
Oggi 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 il machine learning sviluppato in collaborazione dalle principali organizzazioni hardware e software di machine learning. Prima della creazione del progetto OpenXLA, XLA è stato sviluppato all'interno del progetto TensorFlow, ma gli obiettivi fondamentali rimangono gli stessi:
Migliora la velocità di esecuzione. Compila i sottografi per ridurre i tempi di esecuzione delle operazioni di breve durata ed eliminare l'overhead del runtime, creare le operazioni con pipeline per ridurre l'overhead di memoria e specializzare le forme note di tensor 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 combinate automaticamente per eguagliare quelle delle operazioni personalizzate originariamente fuse a mano.
Migliorare la portabilità. Rendi relativamente facile scrivere un nuovo backend per hardware innovativo, in modo che un'ampia parte dei modelli ML possa essere eseguita senza modifiche su quell'hardware. Questo è in contrasto con l'approccio di singole operazioni monolitiche specializzate per il nuovo hardware, che richiede la riscrittura dei modelli per poter utilizzare queste operazioni.
Come funziona
Il compilatore XLA prende i grafici dei modelli dai framework ML definiti in StableHLO e li compila in istruzioni della macchina 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 include questi passaggi:
XLA esegue diversi passaggi di ottimizzazione e analisi integrati sul grafico StableHLO indipendenti dalla destinazione, ad esempio CSE, fusione di operazioni indipendenti dalla destinazione e analisi del buffer per l'allocazione della memoria di runtime per il calcolo. Durante questa 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 con in mente informazioni e esigenze specifiche per il target. Ad esempio, il backend GPU potrebbe eseguire fusioni di operazioni che sono vantaggiose specificamente per il modello di programmazione GPU e determinare come partizionare il calcolo in flussi. In questa fase, i backend possono anche corrispondere a pattern determinate operazioni o combinazioni di queste operazioni con le chiamate alla libreria ottimizzate.
Il backend quindi esegue la generazione di codice specifico per la destinazione. I backend di CPU e GPU inclusi in XLA utilizzano LLVM per infrarossi, ottimizzazione e generazione di codice di basso livello. Questi backend emettono l'IR LLVM necessario per rappresentare il calcolo HLO in modo efficiente, quindi chiamano LLVM per emettere codice nativo da questo IR LLVM.
All'interno di questo processo, il compilatore XLA è modulare, nel senso che è facile eseguire lo slot in un backend alternativo per scegliere come target un'architettura HW nuova. Il backend GPU attualmente supporta le GPU NVIDIA tramite il backend LLVM NVPTX. Il backend della CPU supporta più ISA di CPU.