Questo documento descrive il percorso di un modulo High Level Optimizer (HLO) di XLA dal suo stato iniziale a un eseguibile finale. A volte omettiamo la parola "modulo" e ci riferiamo a questo solo come "HLO".
HLO pre-ottimizzato
Iniziamo con il modulo HLO pre-ottimizzazione. L'HLO pre-ottimizzato non contiene operazioni (ops) considerate interne a XLA, come fusion o bitcast. In questa fase, le operazioni non hanno un layout oppure, se lo hanno, verrà
ignorato. L'HLO pre-ottimizzato viene in genere prodotto da framework di livello superiore
come TensorFlow e JAX. Quando utilizzi il flag XLA -xla_dump_to, l'HLO pre-ottimizzazione viene scaricato in un file con il suffisso del nome file
"before_optimizations.txt".
Ottimizza il modulo HLO
La pipeline XLA:GPU trasforma l'HLO pre-ottimizzato in HLO ottimizzato eseguendo una sequenza di passaggi. I pass possono essere raggruppati semanticamente ed eseguiti nel seguente ordine:
Tessere correlate allo sharding
Sono inclusi pass come Shardy Partitioner o quelli per lo sharding SPMD.
Passaggi di ottimizzazione
Possono essere inclusi sia i pass di legalizzazione sia quelli di semplificazione.
Passaggi di ottimizzazione collettiva
Simile ai passaggi di ottimizzazione, ma si concentra sulle operazioni collettive.
Pass di assegnazione del layout
A ogni operazione HLO viene assegnato un layout che fa parte della forma dell'istruzione. Il layout controlla la disposizione fisica del tensore in memoria.
Esempio di forma con layout:
f32[10,20,30]{2,0,1}
Dopo il tipo di elemento, ci sono le dimensioni logiche della forma, seguite dalla permutazione del layout in ordine crescente. In questo esempio, la dimensione meno importante è 30, la seconda meno importante è 10 e la dimensione più importante è 20.
L'obiettivo dell'assegnazione del layout è ridurre al minimo il numero di trasposizioni fisiche richieste utilizzando una strategia greedy. Inizia con determinati vincoli di layout (ad esempio, le librerie cuDNN/cuBLAS prevedono dimensioni consecutive) e propaga i layout "verso il basso" e poi "verso l'alto" nel grafico HLO. Al termine della propagazione del layout, alcune istruzioni potrebbero avere layout in conflitto, uno propagato da un operando e uno propagato da un utente. Per risolvere questo
conflitto, viene inserita un'istruzione HLO copy che modifica il layout dal
layout dell'operando al layout dell'istruzione.
Passaggi di normalizzazione del layout
Poiché è piuttosto difficile capire la forma fisica, la normalizzazione del layout tenta di riscrivere la forma in modo che utilizzi il layout predefinito {rank-1, rank-2, …, 0}. Nell'esempio precedente, la forma normalizzata sarebbe
f32[20,10,30]{2,1,0}. Le operazioni di copia che modificano i layout vengono riscritte come
combinazione di transpose e bitcast. Dato che al momento non possiamo normalizzare tutte le operazioni, alcune potrebbero ancora avere layout non predefiniti, in particolare gather e dot. Ai confini tra le operazioni normalizzate e
quelle non normalizzate, ci saranno bitcast operazioni che rappresentano una trasposizione, ovvero
una trasposizione con un layout assegnato che la rende fisicamente un'operazione no-op.
La normalizzazione del layout rende esplicite anche alcune trasposizioni implicite, il che è
importante perché codegen può gestire le trasposizioni esplicite con un emettitore
dedicato. Ad esempio, è tecnicamente consentito che una rimodellamento abbia un layout fisico diverso tra operando e risultato (ad es. a causa di un rango diverso). La
passata ReshapeDecomposer che viene eseguita nell'ambito delle passate di normalizzazione del layout
trasforma un'operazione di rimodellamento in una sequenza di transpose, rimodellamento bitcast e
transpose.
Passaggi di ottimizzazione dell'assegnazione del layout dei post
I passaggi più importanti qui sono le fusioni Triton (fusioni GEMM + fusioni Softmax/Layernorm) o le riscritture delle chiamate di libreria. In questo passaggio viene eseguito anche l'ottimizzazione automatica, in cui XLA sceglie tra diversi emettitori, seleziona l'algoritmo migliore per le convoluzioni o i prodotti scalari, trova il miglior tiling per le fusioni gestite dall'emettitore Triton e così via.
Tessere fusion
I due pass principali sono la fusione PriorityFusion e Multi-Output.
In PriorityFusion, formiamo le fusioni guidate dal modello di costo. Durante l'unione
consentiremmo la duplicazione delle operazioni con più utenti se l'operazione può essere unita a tutti
gli utenti. Se possibile, consentiremmo anche di estendere le fusioni Triton Softmax esistenti.
La fusione Multi-Output è un passaggio separato che consente di unire operazioni/fusioni che condividono un operando. Può anche unire operandi/fusioni di operandi in utenti senza
duplicazione aggiungendo output aggiuntivi, in modo che altri utenti dell'operazione da unire possano
essere reindirizzati a questi output. Questo passaggio deve fare attenzione a non introdurre
cicli nel grafico HLO.
Dopo la fusione multi-output, viene eseguita l'eliminazione delle sottoespressioni comuni (passaggio HloCSE), che potrebbe unire nuovamente le operazioni precedentemente duplicate se sono finite nella stessa fusione.
Diverse passate post-fusione
Diversi passaggi relativi ai collettivi (ad esempio la conversione in asincrono o l'applicazione di un determinato ordine relativo dei collettivi).
Infine, eseguiamo CopyInsertion in cui vengono aggiunte copie per garantire che le operazioni
in loco non sovrascrivano i dati ancora necessari altrove.
Al termine dell'ottimizzazione, l'HLO ottimizzato viene scaricato se si utilizza il flag
-xla_dump_to in un file con il suffisso del nome file
"after_optimizations.txt". Se vuoi eseguire il dump di HLO dopo i passaggi intermedi
che modificano effettivamente HloModule, puoi utilizzare il flag
-xla_dump_hlo_pass_re=.* (o un'espressione regolare specifica per limitarlo a
determinati passaggi).
Pianificazione
Un modulo HLO senza una pianificazione ha comunque un certo grado di libertà nell'ordine in cui vengono elaborate le operazioni. Qualsiasi ordinamento topologico che rispetti le relazioni tra operandi/risultati e le dipendenze di controllo è valido. La pianificazione determina l'ordine specifico da utilizzare. Il problema principale in questa fase è il consumo massimo di memoria che dipende dalla durata dei tensori. In un passaggio iniziale, proviamo diversi algoritmi di pianificazione e scegliamo la pianificazione che dovrebbe ridurre al minimo il picco di consumo di memoria. Tieni presente che a questo punto non lavoriamo ancora con buffer fisici (lo faremo in "Assegnazione buffer ") e simuliamo l'utilizzo della memoria.
Poi, LatencyHidingScheduler esegue passaggi e cerca di massimizzare
la sovrapposizione tra calcolo e comunicazione. Tuttavia, ciò potrebbe aumentare nuovamente l'utilizzo della memoria.
Infine, nel caso in cui il consumo massimo di memoria sia superiore alla quantità di memoria disponibile, eseguiamo HloRematerialization. Questo passaggio tenta di ridurre
l'utilizzo di memoria a scapito delle prestazioni, ad esempio alcune fusioni potrebbero essere suddivise e
alcune operazioni potrebbero essere duplicate per avere durate del buffer più brevi. Se si verifica la
rematerializzazione, potrebbe essere utile esaminare i modi per ridurre
i requisiti di memoria sul lato del modello (ad esempio, utilizzando batch di dimensioni inferiori).
Assegnazione del buffer
Immediatamente prima di passare a LLVM IR, eseguiamo i passaggi di assegnazione del buffer che assegneranno le sezioni del buffer a ogni istruzione nel grafico HLO. L'assegnazione del buffer viene eseguita in più passaggi:
HloDataflowAnalysisassegnaHloValues(essenzialmente buffer logici) alle istruzioni. Per le operazioni in loco, è possibile riutilizzareHloValuedi un operando. Un'operazione può definire più di unHloValue(ad es. con una forma di risultato a tupla).HloAliasAnalysistenta di combinare i buffer per le operazioni di aliasing e calcola una mappatura daHloValueaHloBuffer.BufferAssignmentcalcola una mappatura diHloBuffersalle sezioni del buffer all'interno di un buffer di grandi dimensioni in modo che la stessa sezione del buffer non venga utilizzata per diversiHloBufferscon durate sovrapposte. Per le operazioni che potrebbero avere alias, è accettabile che ci sia una leggera sovrapposizione (l'ora di fine di unHloBufferpotrebbe coincidere con l'ora di inizio dell'altroHloBuffer). Quando utilizzi il flag-xla_dump_to, alcune informazioni sull'assegnazione del buffer vengono scaricate in un file con il suffisso "after_optimizations-buffer-assignment.txt".
Thunks
Dopo che un grafico HLO viene ottimizzato e pianificato, viene trasformato in una sequenza lineare di chunk per un backend specifico (CPU o GPU).
In XLA, un thunk è un'astrazione di un'unità di lavoro autonoma che il runtime esegue. Potrebbe trattarsi di un lancio del kernel compilato, di un'operazione specifica, di una chiamata di libreria, di un costrutto di flusso di controllo, di una comunicazione collettiva e così via. Una sequenza di thunk rappresenta l'intero eseguibile per un backend specifico.
Thunk Emission
La procedura di conversione di un calcolo HLO pianificato in una sequenza di thunk è chiamata "emissione di thunk". Questa operazione viene gestita da una classe di emettitori dedicata in ogni backend.
Per il backend GPU, questa operazione viene gestita da
IrEmitterUnnested.
EmitHloComputation scorre l'elenco pianificato di istruzioni HLO in un calcolo e le invia a un metodo Emit... specializzato (ad es.
EmitFusion, EmitConvolutionThunk, EmitWhile). Ciascuno di questi metodi
costruisce gli oggetti Thunk appropriati e li aggiunge alla
sequenza di thunk.
Per il backend della CPU,
ThunkEmitter
svolge questo ruolo ed è organizzato in modo simile. La ThunkSequence finale
è incorporata in CpuExecutable.
Tieni presente che ogni istruzione nel calcolo della voce di un modulo HLO potrebbe
corrispondere a zero (kTuple, kConstant, ..), uno o più (ad esempio
l'istruzione di ordinamento) thunk nella sequenza di thunk finale.
Buffer dei comandi: ottimizzazione dell'esecuzione sulla GPU
L'hardware GPU moderno consente di registrare una sequenza di operazioni GPU (avvii del kernel, copie di memoria e così via) una sola volta e poi riprodurla più volte con un sovraccarico della CPU minimo. Si tratta di un'ottimizzazione delle prestazioni fondamentale, soprattutto per i carichi di lavoro con molti kernel piccoli e a esecuzione rapida. XLA utilizza Command Buffer come astrazione di CUDA Graphs o HIP Graphs. L'interfaccia principale è definita in GpuCommandBuffer.
Un buffer dei comandi è rappresentato in una sequenza thunk da CommandBufferThunk.
L'emettitore non produce questo thunk direttamente dalle istruzioni HLO. Questa operazione viene invece eseguita da CommandBufferConversionPass che viene eseguita sulla sequenza Thunk stessa.
La passata identifica sottosequenze contigue di thunk compatibili (ad es. una
serie di KernelThunk e GemmThunk). Sostituisce quindi la sottosequenza
trovata con un singolo CommandBufferThunk. Il nuovo thunk incapsula la
logica dei thunk originali come elenco di oggetti CommandBufferCmd leggeri.
Quando un CommandBufferThunk viene eseguito per la prima volta su un determinato stream GPU,
la sua sequenza di comandi viene "registrata" in un buffer di comandi hardware. In tutte le
esecuzioni successive, viene emesso un singolo comando alla GPU per "riprodurre"
la sequenza registrata. In questo modo si evita l'overhead della CPU dovuto all'avvio di ogni singolo
kernel.
Eseguibili
Il prodotto finale della pipeline di compilazione XLA è un eseguibile autonomo e specifico per la piattaforma. Questo oggetto contiene tutte le informazioni necessarie per eseguire il programma compilato su un dispositivo di destinazione, ad esempio una CPU o una GPU. È il ponte tra il compilatore e il runtime. I runtime moderni come PJRT utilizzano astrazioni di livello leggermente superiore (vedi PjRtExecutable), ma alla fine eseguono il wrapping di un eseguibile specifico per il backend.
Un Executable contiene diverse informazioni chiave generate
durante la compilazione. Sebbene i contenuti esatti varino a seconda del backend, in genere
includono:
Codice compilato: si tratta del codice macchina di basso livello che verrà eseguito sul dispositivo. Per le CPU, in genere si tratta di uno o più file oggetto. Per le GPU, questo è il codice del dispositivo compilato in formato PTX o HSACO, che viene caricato sulla GPU in fase di runtime.
Piano di esecuzione (ThunkSequence): il fulcro della logica di runtime. Si tratta di una sequenza lineare di oggetti Thunk. Ogni thunk rappresenta una singola unità di lavoro, ad esempio l'avvio di un kernel, la chiamata di una funzione di libreria (ad es. cuBLAS) o la gestione del flusso di controllo. Il runtime esegue il programma iterando questa sequenza.
Layout di memoria (BufferAssignment): questo metadato fondamentale, prodotto da BufferAssigner, descrive il layout di memoria completo per il calcolo. Specifica le dimensioni di ogni buffer e il modo in cui la memoria viene allocata e riutilizzata per parametri, output e valori temporanei. Il runtime lo utilizza per allocare la memoria del dispositivo e passare i puntatori corretti a ogni thunk.
(facoltativo) Modulo HLO: per il debug e la profilazione, l'eseguibile spesso mantiene un riferimento all'HloModule finale ottimizzato da cui è stato compilato.
La creazione dell'eseguibile finale è orchestrata dal compilatore per ogni
backend specifico. Il metodo RunBackend di un'implementazione del compilatore è
il passaggio finale del processo di compilazione, che raggruppa tutti gli artefatti compilati
in un oggetto eseguibile.
GpuCompiler
e
CpuCompiler
hanno come target rispettivamente GPU e CPU.
Quando un utente chiama Execute... su un eseguibile, il runtime utilizza
BufferAssignment per allocare la memoria, quindi richiama ThunkSequence per
avviare le operazioni sul dispositivo utilizzando il codice compilato.