Moltiplicazione di matrici su GPU

Moltiplicazione efficiente di matrici su GPU

Come raggiungere prestazioni di moltiplicazione di matrice di ultima generazione in CUDA.

Questo blog nasce da una improvvisa realizzazione di quanto poco sapessi su come funziona la moltiplicazione di matrici sulla GPU. Avendo fatto così tanti progetti di ML, sento di dover capire come funziona l’operazione più importante in ML: Che cos’è questa cosa chiamata “Tensor Core”? Perché tutti dicono che “il movimento dei dati è il collo di bottiglia”? Quanto veloci possono essere le GPU effettivamente?

Per rispondere a queste domande, ho deciso che devo uscire dalla mia bolla di PyTorch e avventurarmi nell’abisso di CUDA. Ho scritto questo blog per documentare tutto ciò che ho imparato, e spero che chiunque lo legga non debba passare attraverso il dolore di cercare tra la documentazione e il codice di CUDA come ho fatto io.

Se c’è qualcosa che ho imparato in questo viaggio, è che la moltiplicazione di matrici concorrente è DIFFICILE. Una moltiplicazione di matrici efficiente dipende molto dall’hardware specifico che stai usando e dalle dimensioni del problema che stai cercando di risolvere. Non esiste una soluzione universale.

Basta lagnarsi, iniziamo!

• Algoritmo di Dijkstra ponderato dal tempo di percorrenza nelle reti OSM
• Come estendere i DataFrames di Pandas con metodi personalizzati per aumentare la funzionalità e la leggibilità del codice
• Codifica Byte-Pair per principianti

Riepilogo sull’architettura della GPU

Ricordiamoci come funzionano le (NVIDIA) GPU. Una GPU raggiunge il parallelismo eseguendo molti thread. Ogni thread viene eseguito su un singolo CUDA core, anche se in un dato momento, solo un sottoinsieme dei thread è attivo, quindi possono esserci molti più thread rispetto ai core CUDA disponibili. Ogni thread, che sia attivo o meno, ha il proprio set di registri.

Un gruppo di 32 thread è noto come un warp. Tutti i thread in un warp devono essere eseguiti insieme (o essere inattivi insieme). Nella maggior parte dei casi, ci sono molti più warp inattivi che warp attivi, e il warp scheduler è responsabile della scelta di quali warp eseguire in un dato momento. Ciò consente alla GPU di nascondere la latenza degli accessi alla memoria pianificando l’esecuzione di altri warp mentre un warp attende i dati.

Un gruppo di warp è noto come un threadblock. Tutti i warp in un threadblock vengono eseguiti nello stesso Streaming Multiprocessor (SM). Ogni threadblock ha la propria memoria condivisa a cui possono accedere tutti i thread nel threadblock.

Nota: Architetture più recenti

A partire dall’architettura Volta, ogni thread ha anche il proprio program counter e stack di chiamate, ecc. Ciò significa che ogni thread in un warp può eseguire istruzioni diverse contemporaneamente.

L’architettura Volta ha anche introdotto i Tensor Core che sono specializzati per risolvere moltiplicazioni di matrici di dimensioni specifiche. Ogni warp attivo ha accesso a un Tensor Core.

Nell’ultima architettura Hopper, c’è un concetto di cluster di threadblock che rappresenta un gruppo di threadblocks. Questo dà all’utente un controllo più fine sulla pianificazione dei threadblock e consente alla memoria condivisa di un threadblock di essere accessibile da altri threadblocks nello stesso cluster.

Parallelizzare la moltiplicazione di matrici

Supponiamo di voler calcolare:

Diciamo che le dimensioni del problema sono (M, N, K) in questo caso. Per parallelizzare questa operazione, possiamo suddividere A e B in matrici più piccole, moltiplicarle singolarmente e concatenare i risultati per formare C.

In particolare, possiamo partizionare A per righe (ossia M in parti di dimensione M’) e B per colonne (ossia N in parti di dimensione N’) per ottenere:

Possiamo vedere che ogni sotto-matrice in C è indipendente l’una dall’altra, quindi possiamo facilmente parallelizzare il calcolo di ogni sotto-matrice.

Nella pratica, K potrebbe essere troppo grande per caricarlo direttamente in memoria e calcolarci sopra. Invece, un’implementazione tipica dividerà anche K in blocchi di dimensione K’, itererà su ogni blocco e accumulerà (sommando) i risultati parziali. Questo è noto come riduzione seriale-K (a differenza della riduzione parallela-K, vedi la sezione sottostante). Matematicamente, si presenta così:

Nota: Padding

In ogni punto in cui la dimensione del problema non è divisibile per la dimensione della partizione, dobbiamo aggiungere padding. Questo è tipicamente fatto implicitamente quando carichiamo gli input partizionati (𝐴ᵢ,ₖ e 𝐵ₖ,ⱼ) nella memoria di livello inferiore in cui assicuriamo che la partizione caricata (di dimensione M’×K’ per 𝐴ᵢ,ₖ e K’×N’ per 𝐵ₖ,ⱼ) sia sempre “piena” aggiungendo zeri. Si deve prestare particolare attenzione quando si scrivono i risultati di nuovo nella memoria globale per evitare errori fuori dai limiti.

A livello più alto, tre partizioni nidificate consentono di parallelizzare la moltiplicazione di matrici sulla GPU:1. La prima partizione avviene a livello di blocco di thread. Ogni blocco di thread è responsabile del calcolo di Cᵢ,ⱼ = Aᵢ Bⱼ.2. La seconda partizione avviene a livello di warp. Il problema a livello di blocco di thread Cᵢ,ⱼ viene ulteriormente partizionato in modo che ogni warp sia responsabile del calcolo di Cᵢ,ⱼ⁽ᵐⁿ⁾ = Aᵢ⁽ᵐ⁾ Bⱼ⁽ⁿ⁾.3. La terza partizione avviene a livello di istruzione. Alcune istruzioni richiedono input di dimensioni particolari. Ad esempio, le Tensor Core di seconda generazione operano su problemi di dimensione (16, 8, 8) per fp16, mentre un’implementazione diretta su core CUDA tramite moltiplicazione scalare opererebbe semplicemente su dimensione (1, 1, 1). Il problema a livello di warp viene quindi ulteriormente partizionato in modo che ogni chunk abbia una dimensione adatta per l’istruzione: Cᵢ,ⱼ⁽ᵐⁿ⁾⁽ᵃᵇ⁾ = Aᵢ⁽ᵐ⁾⁽ᵃ⁾ Bⱼ⁽ⁿ⁾⁽ᵇ⁾.

Ci sono buone ragioni per cui abbiamo bisogno di tre livelli di partizione, come vedremo nella prossima sezione.

Ridondanza dei dati

La moltiplicazione di matrici può facilmente diventare legata alla memoria se riprendiamo in modo ingenuo i dati dalla memoria globale ai registri ogni volta che effettuiamo un calcolo. La chiave osservazione è che molte delle sotto-istanze Aᵢ e Bⱼ vengono riutilizzate in diverse moltiplicazioni di sub-matrici. Ad esempio, Aᵢ è richiesta per Cᵢ,₁ , Cᵢ,₂ , … e Bⱼ è richiesta per C₁,ⱼ , C₂,ⱼ , … . Possiamo ottenere la migliore efficienza se riusciamo a minimizzare gli spostamenti ridondanti dei dati e riutilizzare i dati caricati il più possibile.

In CUDA, ci sono tre tipi di memoria accessibile dall’utente:

Ecco una visione generale di come viene utilizzato ogni tipo di memoria:

1. Ogni blocco di thread caricherà prima i suoi input richiesti dalla memoria globale nella memoria condivisa. Gli accessi successivi a quei dati verranno quindi gestiti dalla memoria condivisa anziché dalla più lenta memoria globale.
2. All’interno di ogni blocco di thread, ogni warp caricherà prima i suoi input richiesti dalla memoria condivisa nei registri. Gli accessi successivi a quei dati verranno gestiti direttamente dai registri veloci.

Approfondimento dei dettagli

Livello di blocco di thread

Al livello di blocco di thread, il problema viene suddiviso in sottoproblemi di dimensione (M’, N’, K’). Quindi, ogni blocco di thread è responsabile per il calcolo di un frammento di C, indicato come:

Il movimento dati ridondante è minimizzato caricando i sottoparametri di input Aᵢ,ₖ e Bₖ,ⱼ nella memoria condivisa. Quando abbiamo terminato di calcolare Aᵢ,ₖ e Bₖ,ⱼ, il frammento successivo (Aᵢ,ₖ₊₁ e Bₖ₊₁,ⱼ) sarà caricato.

Livello di warp

Al livello di warp, il sottoproblema è ulteriormente suddiviso in sottosottoproblemi di dimensione (M’‘, N’‘, K’’). Quindi, ogni warp è responsabile per il calcolo di un frammento di Cᵢ,ⱼ, indicato come Cᵢ,ⱼ⁽ᵐ ⁿ⁾:

Il movimento dati ridondante è minimizzato caricando i sottosottoparametri di input Aᵢ,ₖ⁽ᵐ ˡ⁾ e Bₖ,ⱼ⁽ˡ ⁿ⁾ nei registri. Qualsiasi accesso a Aᵢ,ₖ⁽ᵐ ˡ⁾ e Bₖ,ⱼ⁽ˡ ⁿ⁾ all’interno di un warp sarà quindi gestito dai registri veloci.

Nota: Distribuzione dei dati tra i registri

Vale la pena notare che i registri sono specifici per il livello di thread. Ciò significa che gli input in un registro non possono essere accessibili da altri thread in un warp. La modalità esatta di come Aᵢ,ₖ⁽ᵐ ˡ⁾ e Bₖ,ⱼ⁽ˡ ⁿ⁾ vengono suddivisi nei registri di ogni thread dipende dall’istruzione specifica utilizzata. La documentazione di NVIDIA su istruzioni di moltiplicazione-accumulo matriciale a livello di warp fornisce una descrizione dettagliata per ogni istruzione.

Livello di Tensor core

Per effettuare effettivamente la moltiplicazione tra matrici, utilizziamo i Tensor Core sulla GPU. La mia GPU (RTX 2060) ha i Tensor Core di seconda generazione, che sono specializzati per risolvere problemi fp16 di dimensione (M’’‘, N’’‘, K’’’) = (16, 8, 8). Quindi, suddividiamo ulteriormente Cᵢ,ⱼ⁽ᵐ ⁿ⁾ per adattarli alla dimensione richiesta dall’istruzione:

In questo caso, tutti gli input sono già nei registri e quindi il costo del movimento dei dati è minimo.

Nota: Tensor Core

Le operazioni di Tensor Core sono istruzioni di livello di warp, il che significa che tutti i thread in un warp devono eseguire contemporaneamente l’istruzione di Tensor Core, collaborando per preparare i dati da consumare da un solo Tensor Core.

Scelta delle dimensioni delle partizioni

Quindi, dato che vogliamo ridurre al minimo il movimento dei dati, dovremmo semplicemente scegliere una dimensione della partizione il più grande possibile per utilizzare tutta la memoria condivisa e i registri, giusto? Beh, non esattamente.

Dimensione della partizione del blocco di thread

Asintoticamente, all’aumentare delle dimensioni del problema, sì, vogliamo utilizzare il più possibile la memoria condivisa e i registri. Tuttavia, per dimensioni di problema piccole, potremmo incontrare due problemi:

1. Una dimensione della partizione grande significa che avremo meno blocchi di thread. Poiché ogni blocco di thread può essere eseguito solo su una SM, potrebbe significare che non possiamo utilizzare tutte le SM.
2. Per dimensioni di problema che non sono divisibili per la dimensione della partizione, dovremo aggiungere più padding agli input. Ciò riduce l’efficienza in quanto verranno effettuati meno calcoli su input significativi.

Un’implementazione tipica potrebbe utilizzare una dimensione di partizione di (M’, N’, K’) = (128, 256, 32).

Dimensione di partizione di warp

In generale, avere una dimensione di partizione di warp grande significa che ci sarà meno spostamento di dati ridondante, ma al costo di avere meno warps. Avere troppi pochi warps significa che non saremo in grado di nascondere la latenza degli accessi alla memoria (perché potremmo finire gli altri warps da programmare mentre il warp corrente attende i dati).

Un’implementazione tipica potrebbe utilizzare una dimensione di partizione di (M”’, N”’, K”’) = (64, 64, 32).

Dimensione della partizione delle istruzioni

Questo è completamente determinato dalle istruzioni supportate dalla tua GPU. Per la mia RTX 2060, l’istruzione ptx per la moltiplicazione di matrici di Tensor Core fp16 (con accumulo fp16) è mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16, che si aspetta input di dimensione (16, 8, 8).

Ancora più ottimizzazioni

Le tecniche sopra descritte possono avvicinarci alle prestazioni teoriche di picco della GPU quando la dimensione del problema è grande. Tuttavia, per dimensioni di problema più piccole, non sono efficienti. Ci sono due tecniche comuni per migliorare ulteriormente le prestazioni della moltiplicazione di matrici: riduzione parallel-K e software pipeline.

Riduzione parallel-K

Nei casi in cui M e N sono piccole, potremmo avere solo alcuni blocchi di thread. Ad esempio, se (M’, N’) = (128, 256) e la dimensione originale del problema ha M ≤ 128 e N ≤ 256, avremo solo un blocco di thread e quindi utilizzeremo solo una frazione della potenza di calcolo della GPU! (Ad esempio, la mia RTX 2060 ha 30 SM, quindi per massimizzare l’utilizzo vogliamo almeno 30 blocchi di thread.)

Nei casi in cui K è grande (anche se M e N sono piccole), possiamo sfruttare maggiore parallelismo mediante la riduzione parallel-K. Ricordiamo che nella riduzione serial-K, ogni blocco di thread itera sulla seguente somma:

e accumula i risultati intermedi in Cᵢ,ⱼ. Nella riduzione parallel-K, assegniamo invece a ogni blocco di thread di calcolare solo un elemento della somma (cioè Aᵢ,ₖ Bₖ,ⱼ). Questo ci consente di aumentare il numero di blocchi di thread in un fattore di K/K’, utilizzando così più SM.

L’aspetto negativo è che ora dobbiamo allocare più memoria per archiviare i risultati di ogni blocco di thread e invocare un secondo kernel per eseguire una riduzione finale sui risultati parziali per ottenere Cᵢ,ⱼ.

Software pipeline

Di solito, CUDA nasconde la latenza degli accessi alla memoria programmando altri warps per eseguire mentre un warp sta aspettando i dati. Ciò richiede che ci siano abbastanza warps per mascherare la latenza.

Tuttavia, il numero di warps è tipicamente relativamente piccolo durante la moltiplicazione di matrici. Questo perché il numero di warps è limitato da “Il numero di registri disponibili per blocco di thread diviso per il numero di registri richiesti per warp”. Per la moltiplicazione di matrici, utilizziamo molti registri per warp per massimizzare il riutilizzo dei dati. Di conseguenza, potremmo non avere abbastanza warps per mascherare la latenza.

“Gli elementi dell’accumulatore occupano tipicamente almeno la metà del budget totale dei registri di un thread”. – Documentazione di CUTLASS

Per attenuare questo effetto, possiamo utilizzare la software pipeline. In sostanza, possiamo (manualmente) pre-caricare in modo asincrono gli input per la prossima iterazione del loop utilizzando istruzioni speciali. Mentre gli input vengono caricati, possiamo continuare a calcolare sull’iterazione corrente. Viene riassunto dal seguente diagramma:

Tutto ciò è reso possibile dal fatto che la GPU è simile a qualsiasi CPU moderna: può pipelinare gli accessi alla memoria e le operazioni aritmetiche fintanto che non c’è una dipendenza dati tra di loro. Ciò è noto come parallelismo a livello di istruzioni.

La moltiplicazione delle matrici in azione

Se vuoi vedere come tutti questi concetti si uniscono in un’implementazione reale, dai un’occhiata alla mia implementazione dell’addestramento di MNIST da zero con CUDA. Lì, ho addestrato un perceptron multistrato su MNIST utilizzando CUDA, ottenendo un aumento di velocità del 6x rispetto a PyTorch ottimizzato per una dimensione nascosta di 128:

GitHub – andylolu2/cuda-mnist

Contribuisci allo sviluppo di andylolu2/cuda-mnist creando un account su GitHub.

github.com

Riferimenti

1. Documenti di CUTLASS2. Documenti di CUDA3. Esempi di CUTLASS