Questa pagina è stata tradotta dall'API Cloud Translation.
Switch to English

Codice MLIR per XLA

XLA opera su HloInstruction ed esegue molte ottimizzazioni su questa rappresentazione, condividendone molte tra dispositivi target. Come alcuni punti viene calcolata una pianificazione lineare e il buffer di memoria viene assegnato staticamente a ciascun valore. Il codegen specifico del dispositivo funziona attraversando questa sequenza e chiamando "emettitori" per generare una rappresentazione adatta al dispositivo (ad esempio una singola funzione LLVM per il calcolo XLA sulla CPU o una sequenza di "thunk" che incapsulano le operazioni della GPU e possibilmente generato PTX quando targeting GPU).

Come fase di staging, al momento stiamo intercettando il processo subito dopo che XLA ha completato la fase di assegnazione del buffer ed emette invece un modulo MLIR nel dialetto lhlo . Da lì eseguiamo il codegen utilizzando componenti MLIR (principalmente dialetto Linalg, affine e GPU) a seconda del dispositivo.

Di seguito è riportato il piano di registrazione per migrare in modo incrementale XLA / GPU utilizzando lhlo come input lhlo .

Compiti

Ospite Dispositivo
Formato di input HloInstruction * (Attività 1) HloInstruction * (Attività 1)
Formato di output xla :: Thunk (attività 2) LLVM IR (Task 3)
  • L'attività 1 modifica il formato di input sia dell'host che del dispositivo da HloInstruction * a LHLO.
  • L'attività 2 modifica il formato di output dell'host dai thunk in "una piattaforma di atterraggio per host" (vedi sotto).
  • L'attività 3 migra l'output del dispositivo da LLVM IR a una qualche forma di MLIR. È facoltativo per questo progetto e vedere la sezione "Migrazione del dispositivo LLVM IR" per i dettagli.

Questo progetto ha la priorità di avere modelli eseguibili end-to-end con emettitori LHLO abilitati il ​​più possibile. Ciò implica che il seguente elenco di obiettivi per priorità:

  • Rendi eseguibile XLA / GPU con emettitori LHLO, con thunk ed emettitori esistenti non modificati.
  • Elimina i riferimenti a HloInstruction * in LHLO, caso per caso:
    • Passare un emettitore legacy a un emettitore basato su MLIR (ad esempio Linalg) o
    • Traduci meccanicamente l'emettitore esistente per acquisire la rappresentazione MLIR (passa allo Standard con GPU Dialect).

Migrazione di thunk (attività 2)

xla :: gpu :: Thunk è una struttura di dati che:

  • Può essere chiamato dall'host (xla :: gpu :: Thunk :: ExecuteOnStream ()).
  • Trasporta vari dati nelle sue sottoclassi.
  • Interagisce con BufferAllocation :: Slice ed StreamExecutor.
  • Lancia i kernel
  • Chiamate in tutte le librerie di runtime.

Il costo di ciò include:

  • Rappresentazione di dati di configurazione specifici dell'operazione (ad es. Configurazioni di convoluzione).
  • Migrazione di forme op e forme di operandi.
  • Rappresenta un albero di thunk (mentre, condizione, ecc.).

Il lavoro di migrazione è indipendente dalla migrazione di LHLO / emettitore. Con risorse limitate, ha la priorità rispetto alla migrazione LHLO / emettitore.

Abbiamo diverse opzioni su come abbassare la parte lato host da LHLO:

  • TFRT
    • (Pro) grandi involucri CUDA e HIP per l'uso.
    • (Pro) facile da implementare chiamate in libreria (cuDNN, cuBLAS, cuFFT, ecc.), Poiché le operazioni TFRT sono interpretate dal codice C ++.
    • (Con) lato host è in fase di sviluppo e non testato.
  • Codice CPU jitted
    • (Pro) grande abilità inferiore. Crea alcuni loop e condizioni ed è fatto.
    • (Con) GPUDialect non modella ancora catene / flussi / asincronicità / allocazione dei dispositivi.
    • (Con) Il supporto del runtime CUDA / HIP è minimo (percorso del toolkit, versione, caricamento dinamico, ecc.).
  • Runtime XLA esistente (interpretazione)

Decisione: adottare TFRT, ma supportare anche il jitting del codice CPU in TFRT.

Migrazione del dispositivo LLVM IR (attività 3)

Un emettitore elementale genera un'operazione target riempiendola elemento per elemento. Ogni elemento di output dipende da un insieme di elementi dagli operandi. Tutti gli elementi sono descritti combinando il buffer con indici dinamici. È sufficiente descrivere quasi tutte le operazioni "matematiche", ma per motivi di prestazioni solo un ampio sottoinsieme di operazioni "matematiche" viene implementato direttamente in (Cpu | Gpu) ElementalIrEmitter.

ElementalIrEmitter è unico in quanto:

  • Una gran parte del codice è condivisa tra XLA / GPU e CPU.
  • Rappresenta una grande porzione di operazioni viste nei modelli, comprese tutte le operazioni basate sugli elementi.
  • La maggior parte delle fusioni dipende esclusivamente da ElementalIrEmitter.
  • È strutturalmente semplice, in quanto descrive un DAG di dipendenza dati tra elementi op ed elementi operando.
  • È per lo più portatile e di alto livello (ad esempio a differenza di GPU kReduce e GPU kCopy).
  • Il supporto dinamico della forma è facile per le operazioni almeno in termini di elementi.

Ora, per tutte le operazioni, emesse in modo elementare o meno, esistono diverse varianti dello stato finale di ciascuna operazione XLA:

  1. Il codice del dispositivo rimane come LLVM IR.
  2. Rifattorizza il vecchio emettitore come LHLO -> MLIR LLVM Dialect:
    • (Costo) Sarà un lavoro a eliminazione se vogliamo infine migrare allo Standard.
    • (Vantaggio) È facile e meccanico. Può essere fatto in un breve periodo.
    • (Vantaggio) Non beneficia di più rispetto a (1).
  3. Riflettono i vecchi emettitori come LHLO -> MLIR GPU + Standard + Loop:
    • (Costo) Il sollevamento degli emettitori esistenti allo standard introduce alcune sfide. Puntatori e GEP devono essere convertiti in MemRef e SubView. Garantire la completezza di amdgpu è un altro.
    • (Costo) XLA / GPU si basa fortemente sui metadati LLVM:
      • range per indici blocco / filo.
      • align , dereferenceable , invariant.load , alias.scope , noalias per load / stores.
      • llvm.loop.unroll.disable , llvm.loop.unroll.full , llvm.loop.vectorize.enable per i cicli sequenziali.
    • (Vantaggio) Può essere a lungo termine. Più portatile.
  4. Trasforma i vecchi emettitori in LHLO -> Linalg e scrivi nuovi emettitori Linalg
    • (Costo) Questo caso per caso. Rispetto alle opzioni precedenti, una nuova implementazione che corrisponde alle prestazioni di XLA deve passare attraverso il benchmark <-> ottimizzare il flusso di lavoro, che può essere un costo significativo per alcune operazioni.
    • (Vantaggio) stack unificato; supporto della comunità; portabilità; più potenziali di ottimizzazione.

conclusioni:

  • Non andare per (2). (1) o (3) sono semplicemente migliori di (2). (2) costa più di (1), poiché richiede molti refactoring meccanici. Con (1) possiamo ancora raggiungere l'obiettivo di consentire a XLA di raccogliere gli emettitori MLIR. Questo è facendo LHLO -> LLVM IR -> esegui emettitori di dispositivi legacy.
  • Le operazioni di ElementalIrEmitter vanno per (4), ma non in modo incrementale. Non c'è modo di farlo op per op, perché tutte le op emesse elementalmente sono collegate nello stesso grafico. Questo lavoro può anche servire come punto di unificazione di diverse forze in corso (xla / service / mlir_gpu, il generatore del kernel, Linalg).
  • Tutte le altre operazioni vanno per (1). Come obiettivo tratto, potrebbero essere migrati a (3) o (4).

Definizione delle priorità

Mentre tutte e tre le attività sopra menzionate sono parallelizzabili, con risorse limitate devono essere serializzate. La definizione delle priorità si concentra su risultati visibili per il completamento di ogni attività.

La priorità è: Task1 (LHLO per emittenti legacy)> Task 2 (Thunks)> Task 3 (emettitori MLIR).

Entro la fine dell'attività 1, gli utenti di XLA possono generare un LHLO (ad es. Generatore di kernel) ed eseguirli. Il formato della compilation non sarà MLIR serializzabile.

Alla fine dell'attività 2, LHLO si riduce a MLIR corretto, serializzabile. Ciò consente la compilazione offline.

Entro la fine dell'attività 3, tutti gli emettitori XLA sono basati su MLIR nella sua implementazione.

Design dettagliato

Passaggio 1: (Attività 1) Completa LHLO e crea legacy Emitters Take LHLO

Questo passaggio consente a tutti gli emettitori XLA / GPU esistenti di interagire con le operazioni MLIR. Questo passaggio è puro refactoring e NFC.

Questo passaggio è principalmente meccanico, ma vale la pena notare le seguenti discrepanze tra un HloComputation non desiderato e LHLO:

  • Ogni HloInstruction ha accesso diretto ai suoi operandi (un DAG del flusso di dati). Al contrario, ogni operazione LHLO ha accesso solo ai suoi buffer degli operandi (un bipartito tra operazioni e buffer). Le operazioni di LHLO devono passare attraverso le catene use-def per accedere alle loro operazioni di operando.
  • Gli emittenti legacy non annidati empiricamente non accedono quasi mai ai loro operandi. L'unica eccezione è kReduce.
  • Gli emittenti legacy non nidificati accedono a BufferAssignment solo per ottenere sezioni, non per accedere a strutture di dati aux come dataflow_analysis () o alias_analysis (). llvm_ir crea il proprio alias_analysis () in base alle informazioni sulle sezioni.

La conclusione è che LHLO dovrebbe adattarsi senza problemi importanti.

Passaggio 2: supporto (facoltativo) di profilazione

Questo passaggio è necessario solo se iniziamo a scartare parte della logica XLA Thunk (vedere il passaggio successivo).

Prima di attivare effettivamente qualsiasi emettitore basato su MLIR, è necessario creare un profilo per gli emettitori basati su MLIR.

Attualmente XLA esegue la propria profilazione chiamando il timer di StreamExecutor. Il timer sotto il cofano inserisce due eventi prima e dopo il lancio del kernel e misura il tempo di sincronizzazione tra questi due eventi.

Esistono circa tre approcci per supportare la profilazione in MLIR:

  • Esegui un profiler end-to-end
  • Aggiungi un profilo operativo per ogni operazione in LHLO, usando un profiler iniettato.

L'approccio "end-to-end" è trasparente per MLIR, ma presenta lo stesso problema che rende XLA non lo utilizza in primo luogo: le chiamate in libreria raccolte da un profiler (nvprof / ...) non possono essere facilmente correlate a HLO ops. Ad esempio, cuDNN avvia più kernel per ciascun HLO ed è difficile stabilire quali kernel corrispondono a quale HLO.

L'approccio del "profilatore iniettato" richiede:

  • LHLO prende un profiler come parametro.
  • inserendo profile.start / profile.end prima e dopo ogni op.
  • un passaggio da quel profilo inferiore. {inizio, fine} a un'implementazione C ++.

La profilazione esatta non può essere eseguita facilmente per le operazioni generate da MLIR, poiché:

  • MLIR non ha un timer, né dipende da TFRT / StreamExecutor.
  • MLIR non chiama facilmente le funzioni C con parametri complicati.

Passaggio 3: (Attività 2) Migrazione di thunk

Come nota, ci sono circa tre tipi di thunk:

  • KernelThunk, che avvia un kernel.
  • Thunk di flusso di controllo, che ha la logica del flusso di controllo host (condizionale, mentre, per, sequenza) e avvia i kernel del corpo.
  • Thunk di libreria: cuDNN, cuBLAS, cuFFT, NCCL, ecc.

Il piano è:

  • Rendi Thunks (de) serializzabile.
  • Aiuta a migliorare TFRT in uno stato in cui può supportare queste semantiche.
  • Man mano che lo stato migliora, migra i singoli thunk in modo incrementale.

Questi elementi di azione sono ordinati solo parzialmente. L'effettivo ordine di esecuzione / parallelismo ingegneristico deve essere valutato man mano che procede.

Passaggio 4: (Attività 3) Migrated ElementalIrEmitter

Una volta che la profilazione è pronta, possiamo completare e ottimizzare tutti gli emettitori basati su ElementalIrEmitter in MLIR. Quindi li accendiamo per impostazione predefinita, supponendo che tutti questi emettitori basati su MLIR utilizzino un singolo flusso.

Si noti che è utile migrare anche ElementalIrEmitter XLA / CPU, poiché condividono gran parte del codice.

Con tutto il benchmarking e la performance caccia eseguita (TODO: definire la parità delle prestazioni), attiviamo il nuovo emettitore elementale basato su MLIR ed eliminiamo l'ElementalIrEmitter legacy.

Questo passaggio fornisce anche facili transizioni di fusione (operazioni nidificate) per la migrazione successiva.

Passaggio 5: supporto multi-stream o rilascio

Non possiamo eliminare alcuni degli emettitori fino a quando non lo supportiamo in MLIR o non rilasciamo la funzione. È una quantità relativamente grande di lavoro in MLIR e una piccola quantità di guadagno per XLA. Dovremmo esaminare gli attuali utenti di utenti XLA / GPU multi-stream e provare a eliminare questa funzione se ragionevole.

Passaggio 6: (Operazione 3) Operazioni dispositivo migrate

Questo passaggio migra tutte le operazioni non indesiderate, quindi possiamo eliminare tutti gli emettitori non indesiderati.

Questo richiede una riscrittura / refactor per kCopy e kReduce. kReduce ha già lavorato a sufficienza, quindi resta da vedere l'effettiva quantità di lavoro che deve essere fatto.