Programmare le GPU Apple tramite Go e Metal Shading Language

Programmare le GPU Apple con Go e il linguaggio Metal Shading

Indagare su Go, Cgo, Metal Shading Language, Metal Performance Shaders e confrontare diversi approcci alla moltiplicazione delle matrici

Di seguito descriverò il processo di utilizzo di cgo per l’interfacciamento tra Go e C nativo, come questo può essere utilizzato per interfacciarsi con i binding Objective-C per il framework Metal Performance Shaders di Apple, come interfacciarsi con il codice GPU personalizzato (shader) scritto in Metal Shading Language, e infine eseguire il benchmarking di tutto ciò rispetto alle operazioni di moltiplicazione delle matrici scritte a mano e basate su Go e OpenBLAS. Questo è stato scritto per essere eseguito sul mio MacBook M2.

Il layout della sorgente, disponibile qui su GitHub, si presenta così:

È molto, quindi lo suddividerò in queste sezioni, oppure sentiti libero di passare direttamente ai benchmark.

• GPU e Parallelismo a Punto Fluttuante
• Nozioni di base su Metal GPU
• Metal Shading Language
• Binding Objective-C
• Framework Metal Performance Shaders
• Go e cgo
• Baselines di implementazione Go e OpenBLAS
• Risultati

GPU e Parallelismo a Punto Fluttuante

Assumo che la maggior parte delle persone a questo punto sia intuitivamente familiare con il concetto che le GPU sono incredibilmente potenti per certi tipi di compiti computazionali, specialmente quelli che supportano l’apprendimento automatico. Non è stato fino a quando ho iniziato a sperimentare con Metal che ho capito di persona quanto possano essere più potenti rispetto alle CPU.

• API di Speech del Web Cosa funziona, cosa non funziona e come migliorarlo collegandolo a un modello di linguaggio GPT
• Grandi modelli incontrano Big Data Spark e LLMs in armonia
• Reti neurali grafiche cooperative

Le GPU sono progettate per essere estremamente efficienti nelle operazioni floating-point massivamente parallele che richiedono un’elevata larghezza di banda della memoria. Il mio MacBook M2 ha 8 core CPU e 8 core GPU, ma a titolo di confronto, il Nvidia RTX 4090 contiene 16384 core e l’ H100 contiene 16896 core CUDA con centinaia di core tensoriali specializzati aggiuntivi. Le GPU di solito supportano l’elaborazione SIMD, che consente loro di eseguire contemporaneamente la stessa istruzione su più punti di dati.

Al di fuori della grafica, la moltiplicazione delle matrici e i compiti di algebra lineare in generale beneficiano di questa concorrenza grazie ai loro algoritmi altamente parallelizzabili. Questo supporta a sua volta i carichi di lavoro principali dell’apprendimento automatico come l’addestramento e l’inferenza [1] [2]].

CUDA è probabilmente la piattaforma di programmazione per GPU più conosciuta, specifica per l’hardware Nvidia. Ci sono anche framework matematici disponibili per OpenGL. Framework come TensorFlow e PyTorch possono integrarsi facilmente e ragionevolmente in modo trasparente con l’hardware GPU. Questo è stato un interessante articolo sul miglioramento delle prestazioni dell’integrazione di framework basati su GPU Metal nella libreria NLP spaCy.

Concetti di base delle GPU Metal

Programmare direttamente il calcolo GPU non è semplice come scrivere codice per CPU su dispositivo. Quando si lavora con il framework Metal di Apple, una serie approssimativa di operazioni per eseguire il codice sulla GPU è la seguente:

• Trovare un dispositivo GPU appropriato
• Creare una coda per l’esecuzione dei comandi (ad esempio la MTLCommandQueue)
• Incapsulare i puntatori agli array di dati in un buffer strutturato; se i dati sono codice eseguibile, allora uno stato della pipeline, altrimenti un buffer regolare. Le GPU Apple utilizzano uno spazio di memoria unificato, il che significa che non dobbiamo effettivamente copiare i dati in memoria fisica specifica della GPU
• Eseguire il buffer dei comandi e attendere i risultati o impostare un gestore di eventi alla sua completamento
• Estrarre byte da un buffer di risposta e formattare localmente con il codice del programma CPU

La programmazione GPU grezza utilizza un modello asincrono.

Linguaggio di shading Metal

Il linguaggio di shading Metal è un derivato di C++14 che può essere utilizzato per comporre logiche personalizzate (chiamate “shaders”) da eseguire sulle GPU compatibili con Metal. In generale, e se possibile, probabilmente è meglio utilizzare il framework MPS (discusso in seguito) per funzionalità equivalenti quando possibile: tende ad essere altamente ottimizzato per classi comuni di casi d’uso allineati con GPU (come la moltiplicazione di matrici o reti neurali).

Il debugging del codice MSL è abbastanza difficile. È possibile utilizzare il Shader Debugger attraverso Xcode, ma se si desidera ispezionare o stampare valori intermedi senza Xcode, è necessario scrivere i dati in un buffer di debug di risposta e analizzare le primitive nel wrapper C++ o Objective-C.

Le funzioni MSL sono esposte come interfacce pubbliche tramite la designazione kernel. Il framework Metal passa ID per il contesto del thread di chiamata corrente o del gruppo di thread, che possono essere utilizzati per garantire scritture non sovrapposte. I thread possono essere rappresentati da un sistema di ID tridimensionale; le dimensioni di questo spazio di thread vengono configurate nel codice wrapper.

Di seguito è riportata un’implementazione dell’algoritmo di moltiplicazione di matrici naive, combinato con alcuni loop unrolling che sorprendentemente hanno migliorato significativamente le prestazioni. Questo è solo a scopo di confronto; normalmente la funzionalità MPSMatrixMultiplication di MPS sarebbe più adatta.

kernel void matrix_multiply_naive(device const MatrixParams *params, constant float *A, constant float *B, device float *C, uint2 gid [[thread_position_in_grid]]) {
  if (gid.x >= params->a_rows || gid.y >= params->b_cols) {
    return; // Questo thread è al di fuori delle dimensioni della matrice, non fare nulla
  }
  float sum = 0.0;
  int k;
  // Srotolamento del ciclo; migliora le prestazioni di una grande percentuale
  for (k = 0; k <= params->a_cols - 4; k += 4) {
    sum += A[gid.x * params->a_cols + k] * B[k * params->b_cols + gid.y];
    sum += A[gid.x * params->a_cols + k + 1] * B[(k + 1) * params->b_cols + gid.y];
    sum += A[gid.x * params->a_cols + k + 2] * B[(k + 2) * params->b_cols + gid.y];
    sum += A[gid.x * params->a_cols + k + 3] * B[(k + 3) * params->b_cols + gid.y];
  }
  // Gestisci gli elementi rimanenti
  for (; k < params->a_cols; ++k) {
    sum += A[gid.x * params->a_cols + k] * B[k * params->b_cols + gid.y];
  }
  C[gid.x * params->b_cols + gid.y] = sum;
}

Ho implementato anche una funzione di trasposizione “naive” in MSL per il confronto. Data una matrice trasposta, si tratta di un semplice adattamento della logica sopra, il cui ciclo interno scorre le righe di B anziché le colonne:

// Srotolamento del ciclo; migliora le prestazioni di una grande percentuale
for (k = 0; k <= params->a_cols - 4; k += 4) {
  sum += A[gid.x * params->a_cols + k] * B[gid.y * params->b_cols + k]; // Nota che qui è gid.y * colonne più k
  sum += A[gid.x * params->a_cols + k + 1] * B[gid.y * params->b_cols + k + 1];
  sum += A[gid.x * params->a_cols + k + 2] * B[gid.y * params->b_cols + k + 2];
  sum += A[gid.x * params->a_cols + k + 3] * B[gid.y * params->b_cols + k + 3];
}
// Gestisci gli elementi rimanenti
for (; k < params->a_cols; ++k) {
  sum += A[gid.x * params->a_cols + k] * B[gid.y * params->b_cols + k];
}

Ho discusso questo approccio in un precedente post sul blog come un modo piuttosto semplice per migliorare le prestazioni dello scalar della versione “naive”, almeno sulle CPU. Ne parlerò meglio più avanti.

Collegamenti Objective-C

Il framework Metal fornisce la possibilità di compilare una libreria a partire dal codice sorgente Metal. Una volta caricati i contenuti del file, il codice di collegamento cerca le funzioni kernel per nome e inizializza una nuova MTLComputePipelineState che rappresenta il codice della funzione compilato.

id<MTLDevice> device = MTLCreateSystemDefaultDevice(); // Crea il dispositivo Metal
// Compila e inizializza una nuova libreria a partire dal percorso del codice sorgente
MTLCompileOptions *compileOptions = [MTLCompileOptions new];
compileOptions.languageVersion = MTLLanguageVersion3_0;
NSString *ss = [NSString stringWithUTF8String:source_path];
id<MTLLibrary> lib = [device newLibraryWithSource:ss options:compileOptions error:&error];
// Crea una rappresentazione della funzione kernel di moltiplicazione "naive" nella libreria Metal creata sopra
id<MTLFunction> naiveFunction = [lib newFunctionWithName:@"matrix_multiply_naive"];
// Crea il nuovo stato del pipeline di calcolo
id<MTLComputePipelineState> pipelineStateNaive = [device newComputePipelineStateWithFunction:naiveFunction error:&error];

<!–Per chiamare effettivamente il codice nativo di Metal, la configurazione del thread deve essere impostata, e i buffer GPU devono essere inizializzati.

[computeEncoder setComputePipelineState:pipelineStateNaive];MTLSize threadsPerGrid = MTLSizeMake(params->a_cols, params->a_rows, 1);// Calcola una dimensione del gruppo di thread.// https://developer.apple.com/documentation/metal/calculating_threadgroup_and_grid_sizes?language=objcNSUInteger w = pipelineStateNaive.threadExecutionWidth;NSUInteger h = pipelineStateNaive.maxTotalThreadsPerThreadgroup / w;MTLSize threadsPerThreadgroup = MTLSizeMake(w, h, 1);// Codifica gli input della funzione kernel[computeEncoder setBytes:params length:16 atIndex:0];[computeEncoder setBuffer:bufferA offset:0 atIndex:1];[computeEncoder setBuffer:bufferB offset:0 atIndex:2];[computeEncoder setBuffer:bufferC offset:0 atIndex:3];// Codifica il comando di calcolo.[computeEncoder dispatchThreads:threadsPerGrid   threadsPerThreadgroup:threadsPerThreadgroup];// Termina il passaggio di calcolo.[computeEncoder endEncoding];// Esegue il comando.[commandBuffer commit];

Questo è molto, quindi illustrerò le relazioni qui:

Framework delle Metal Performance Shaders

Il framework MPS è una libreria ad alte prestazioni fornita da Apple per l’utilizzo con la sua famiglia di GPU Metal. Offre funzionalità dalle attività di immagine al supporto di reti neurali.

Le API sono principalmente disponibili tramite Swift o Objective-C, anche se è disponibile anche una libreria Metal-cpp da utilizzare.

L’API MPSMatrixMultiplication è relativamente facile da usare. Come nel codice MSL sopra, i comandi MPS devono ancora essere codificati nel MTLCommandBuffer e commessi asincronamente per l’esecuzione.

// Definisci le "descrizioni" delle matrici, tenendo conto della dimensionalità e delle dimensioni in byte della matriceMPSMatrixDescriptor *descriptorA = [MPSMatrixDescriptor matrixDescriptorWithDimensions:a_rows  columns:a_cols  rowBytes:a_cols * sizeof(float)  dataType:MPSDataTypeFloat32];MPSMatrixDescriptor *descriptorB = [MPSMatrixDescriptor matrixDescriptorWithDimensions:b_rows  columns:b_cols  rowBytes:b_cols * sizeof(float)  dataType:MPSDataTypeFloat32];// Matrice di outputMPSMatrixDescriptor *descriptorC = [MPSMatrixDescriptor matrixDescriptorWithDimensions:a_rows  columns:b_cols  rowBytes:b_cols * sizeof(float)  dataType:MPSDataTypeFloat32];// Inizializza le rappresentazioni delle matrici utilizzando le descrizioni sopra e i buffer delle matriciMPSMatrix *matrixA = [[MPSMatrix alloc] initWithBuffer:bufferA descriptor:descriptorA];MPSMatrix *matrixB = [[MPSMatrix alloc] initWithBuffer:bufferB descriptor:descriptorB];MPSMatrix *matrixC = [[MPSMatrix alloc] initWithBuffer:bufferC descriptor:descriptorC];// Crea l'istanza per la moltiplicazioneMPSMatrixMultiplication *matrixMultiplication = [[MPSMatrixMultiplication alloc] initWithDevice:device  resultRows:a_rows  resultColumns:b_cols  interiorColumns:a_cols];// Codifica il comando di moltiplicazione nel buffer dei comandi per la GPUid<MTLCommandBuffer> commandBuffer = [commandQueue commandBuffer];[matrixMultiplication encodeToCommandBuffer:commandBuffer  leftMatrix:matrixA  rightMatrix:matrixB  resultMatrix:matrixC];

Go e cgo

Non mi piace particolarmente lavorare con Objective-C, e lo scopo di questo programma è eseguire il codice sulla GPU originario di un programma Go.

Cgo è una funzionalità del linguaggio Go che consente al compilatore Go di comprendere le direttive del compilatore contenute nei commenti relativi al codice C nativo. Supporta una versione dell’interfaccia di funzione esterna.

La configurazione delle direttive è un po’ delicata, ma tutti i commenti immediatamente precedenti la riga import "C" (chiamato “il preambolo”) verranno interpretati come importazione di intestazione o argomenti di compilazione durante la compilazione del codice C di riferimento. Ad esempio:

/*#cgo LDFLAGS: -framework Foundation -framework CoreGraphics -framework Metal -framework MetalPerformanceShaders -L/opt/homebrew/opt/openblas/lib -lopenblas#include <stdlib.h>#include "metal.h"*/import "C"

• Passa i flag di collegamento al linker tramite LDFLAGS sulla riga di comando
• Compila il codice C con l’intestazione standard stdlib.h
• Compila il codice C con l’intestazione del progetto locale metal.h

Ho fatto alcuni tentativi ed errori per ottenere il giusto set di flag di collegamento su MacOS.

• Foundation: librerie di base
• CoreGraphics: necessario su MacOS per interfacciarsi con la GPU
• Metal: librerie e supporto del linguaggio per Metal, incluso MSL
• MetalPerformanceShaders: librerie per MPS sopra discusso

Alla fine si è scoperto che Apple include un’implementazione BLAS nel suo framework Accelerate, quindi oltre a installare OpenBLAS tramite brew, è necessario fornire anche la posizione della libreria durante il collegamento:

-L/opt/homebrew/opt/openblas/lib -lopenblas

La direttiva go:embed consente ai programmi Go di includere file durante il tempo di compilazione, che è utile in questo caso quando vogliamo passare i contenuti del file di origine MSL (mm.metal) al framework Metal, come sopra discusso, per la compilazione.

//go:embed mm.metalvar source string// Compila il codice sorgente dello shader e inizializza i pipeline. Il parametro metalSource// contiene i contenuti di un file embedded di Metal Shading Language. Func Compile(metalSource string) { // Incapsula la stringa in una stringa C src := C.CString(metalSource) // Elimina la stringa sopra dopo l'inizializzazione della coda dei comandi defer C.free(unsafe.Pointer(src)) // Compila la sorgente, inizializza i pipeline e la coda dei comandi C.initializePipelineAndCommandQueue(src)}

Le referenziazioni a C sopra interfacciano le API di C tramite cgo, ad esempio:

// Chiama initializeMTLBuffers dai binding Obj-C C.initializeMTLBuffers( a_data, // Puntatore opaco di input per A b_data, // Puntatore opaco di input per B C.int(4), // Converte 4 in tipo intero C C.int(a.Size()), C.int(b.Size()), C.int(a.Rows * b.Cols)) params := MatrixParams{ a_rows: int32(a.Rows), a_cols: int32(a.Cols), b_rows: int32(b.Rows), b_cols: int32(b.Cols),} // Restituisce un puntatore non sicuro a questa struttura MatrixParams, convertito // nella rappresentazione C nativa definita nell'intestazione condivisa return (*C.MatrixParams)(unsafe.Pointer(&params));

Si noti che questo significa che C è una parola chiave riservata e non può essere utilizzata come nome di variabile.

Implementazione di base in Go e OpenBLAS

Volevo confrontare le prestazioni della moltiplicazione di matrici basata sulla GPU con implementazioni di livello superiore, come la libreria Gonum, oltre a implementazioni intuitive scritte a mano (e relativamente inefficienti).

Ho implementato diversi algoritmi in Go, incluso questo algoritmo di trasposizione parallelo naive, che divide in modo ingenuo il lavoro di moltiplicazione tra N goroutine:

func (a Matrix[T]) TransposeMultParallel(b *Matrix[T]) *Matrix[T] { if a.Cols != b.Rows {  panic("matrici di dimensioni errate per la moltiplicazione") } c_data := make([]T, a.Rows*b.Cols) t := b.Transpose() var wg sync.WaitGroup for i := 0; i < a.Rows; i++ {  wg.Add(1) // Aggiunge un conteggio al WaitGroup per la nuova goroutine  go func(i int) { // Avvia la goroutine   defer wg.Done() // Riduce il conteggio quando la goroutine è completata   ptr := i * b.Cols   for j := 0; j < b.Cols; j++ {    var sum T = 0.0    for k := 0; k < a.Cols; k++ {     sum += a.At(i, k) * t.At(j, k)    }    c_data[ptr+j] = sum   }  }(i) } wg.Wait() // Attende che tutte le goroutine siano completate return InitMatrixWithData(a.Rows, b.Cols, c_data)}

Gonum BLAS è una libreria puramente Go che implementa le interfacce BLAS. Tuttavia, può anche essere configurata per deviare le operazioni algebriche a un’implementazione BLAS in codice nativo come OpenBLAS tramite netlib.

Ho mostrato sopra come cgo può essere configurato per collegarsi correttamente a un’installazione di OpenBLAS su MacOS. Nel codice dell’applicazione, è possibile impostare direttamente l’implementazione BLAS preferita. Dal codice di benchmark:

// Converti array primitivi in tipi di matrice densa gonumgonum_a := mat.NewDense(a_rows, a_cols, a64_data)gonum_b := mat.NewDense(b_rows, b_cols, b64_data)gonum_c := mat.NewDense(a_rows, b_cols, nil)gonum_d := mat.NewDense(a_rows, b_cols, nil)// Configura Gonum per utilizzare l'implementazione Go di default di Gonumblas64.Use(gonum.Implementation{})// Esegui una moltiplicazione utilizzando l'implementazione BLAS di Gonumstart = time.Now()gonum_c.Mul(gonum_a, gonum_b)bdata.TimeGonumNative(start)// Configura Gonum per utilizzare Netlib, che inoltra le operazioni a un'implementazione di BLAS in codice C nativo (in questo caso OpenBLAS)blas64.Use(netlib.Implementation{})// Esegui una moltiplicazione utilizzando l'implementazione OpenBLAS tramite l'API di Gonumstart = time.Now()gonum_d.Mul(gonum_a, gonum_b)bdata.TimeGonumOpenBLAS(start)

Risultati

Il mio codice di benchmarking esegue alcuni tentativi di ognuna delle seguenti implementazioni di moltiplicazione di matrici e riporta il tempo medio richiesto per moltiplicare due matrici quadrate di dimensionalità progressivamente crescente:

- Moltiplicazione ingenua, in Go- Moltiplicazione ingenua trasposta, in Go- Moltiplicazione ingenua trasposta parallelizzata con goroutine, in Go- Moltiplicazione BLAS puramente Go-based di Gonum- Moltiplicazione BLAS di OpenBLAS avvolta da Gonum, scritta in C- Moltiplicazione ingenua implementata manualmente, in MSL, su GPU- Moltiplicazione ingenua trasposta implementata manualmente, in MSL, su GPU- Framework Metal Performance Shaders, chiamato da Objective-C, su GPU

L’output del benchmark appare così (i floating point sono in ms):

2023-12-01 11:12:51.644 go-mm[75818:22427382] Utilizzo del dispositivo predefinito Apple M2elementi ingenuo trasposta trasposta_parallela naive_metal metal_transpose mps gonum openblas
160000 196.00 201.00 42.00 8.00 9.67 0.33 4.67 6.00
250000 381.33 387.67 80.67 11.00 11.67 0.00 8.33 21.00
360000 801.00 789.33 159.33 19.00 16.33 0.00 14.33 4.67
490000 1228.00 1075.00 411.00 23.67 24.33 1.00 26.67 16.33...

Un rapido disegno tramite matplotlib

Come ci si potrebbe aspettare, le mie implementazioni scritte a mano di Go sono relativamente fuori controllo. Infatti, gli altri approcci sono così veloci che non si possono nemmeno distinguere nel grafico. Ecco l’istogramma scorrevole dell’utilizzo della GPU durante questa esecuzione.

Puoi vedere che la GPU non è particolarmente occupata perché la maggior parte del tempo viene trascorso nelle operazioni della CPU. Ecco un’altra esecuzione, escludendo le tre tecniche di moltiplicazione più lente:

Acirca 16M di elementi (4k x 4k), Gonum inizia a degradarsi. Puoi vedere chiaramente che le operazioni basate su GPU e OpenBLAS superano le implementazioni Go puri. Guardando solo le approcci basate su GPU:

Ecco un paio di note interessanti:

• La libreria Metal Performance Shaders è incredibilmente veloce
• Non c’è una differenza di prestazioni reale tra gli approcci ingenui e gli approcci ingenui trasposti

Per il secondo punto: ciò non è simile alle caratteristiche di prestazione delle due implementazioni basate su Go sopra. Risulta che i modelli favorevoli di accesso alla cache per le CPU non funzionino allo stesso modo per le GPU e il modo in cui i loro gruppi SIMD (o warps) accedono alla memoria. Guarda qui l’utilizzo della GPU per fare un confronto:

Ora, guardando solo OpenBLAS e MPS – i due approcci più veloci:

Intorno a 35M di elementi, l’implementazione di OpenBLAS inizia a degradarsi, mentre MPS rimane stabile. La differenza qui è piuttosto notevole, con quest’ultimo che completa le stesse operazioni di moltiplicazione tra matrici di 35M di elementi in meno del 15% del tempo. È ragionevole presumere che questa differenza continui a crescere con la cardinalità delle matrici.

Ora, naturalmente, ci sono probabilmente differenze algoritmiche tra questi due approcci, quindi non si tratta di un paragone equo tra CPU e GPU. Se rappresento graficamente le differenze di prestazioni tra le mie due implementazioni scritte a mano, appare così:

Quello che sta dicendo è che l’implementazione ingenua basata su MSL completa la moltiplicazione di 5M elementi in soli il 1% del tempo della mia implementazione in Go, e sembra che questo rapporto si stia migliorando a favore della GPU nel tempo.