Programmazione delle GPU Apple tramite Go e Metal Shading Language |  di Mike Cvet |  Dicembre 2023

 | Intelligenza-Artificiale

Studio di Go, Cgo, Metal Shading Language, Metal Performance Shader e analisi comparativa di diversi approcci alla moltiplicazione delle matrici

fotografato da Etienne Martin SU Unsplash

Di seguito descriverò il processo di utilizzo cgo per interfacciare tra Go e C nativo, come questo può essere utilizzato per interfacciarsi con i collegamenti Objective-C per Apple Shader per prestazioni in metallo framework, come interfacciarsi con costume Codice GPU (shader) scritto nel file Linguaggio di ombreggiatura dei metallie infine confrontando tutto ciò con quello scritto a mano e OpenBLAS Operazioni di moltiplicazione di matrici basate su Go. Questo è stato scritto per funzionare sul mio MacBook M2.

La disposizione della fonte, disponibile qui su GitHubSomiglia a questo:

Layout di alto livello di concetti, tipi e hardware all'interno del wrapper Objective-C
Codice sorgente, libreria e layout del dispositivo di alto livello

È molto, quindi lo suddividerò qui in queste sezioni o sentiti libero di saltare semplicemente diritto ai parametri di riferimento.

Presumo che la maggior parte delle persone a questo punto abbia intuitivamente familiarità con il concetto che le GPU sono incredibilmente potenti in determinati tipi di attività computazionali; soprattutto alcuni che supportano l’apprendimento automatico. È stato solo quando ho iniziato a giocare con il Metal che ho capito in prima persona come farlo tanto possono essere più potenti delle CPU.

Le GPU sono estremamente efficienti nelle operazioni a virgola mobile massivamente parallele che richiedono un’elevata larghezza di banda della memoria. Il mio MacBook M2 ha 8 core CPU e 8 core GPU, ma per confronto, il Nvidia RTX 4090 contiene 16384 nuclei e il H100 contiene 16896 CUDA Core con centinaia di tensor core specializzati aggiuntivi. Le GPU solitamente supportano SIMD elaborazione, consentendo loro di eseguire la stessa istruzione simultaneamente su più punti dati.

Al di fuori della grafica, la moltiplicazione di matrici e i compiti algebrici lineari in generale beneficiano di questa concorrenza grazie ai loro algoritmi altamente parallelizzabili. Questo a sua volta supporta i carichi di lavoro principali del machine learning come training e inferenza (1) (2).

CUDA è probabilmente la piattaforma di programmazione GPU più conosciuta, specifica per l’hardware Nvidia. Sono disponibili anche strutture matematiche per OpenGL. Framework come TensorFlow e PyTorch possono integrare facilmente e ragionevolmente trasparente con l’hardware GPU. Questo è stato un articolo interessante sui miglioramenti delle prestazioni derivanti dal supporto dei framework GPU basati su Metal nel libreria PNL spaCy.

Programmare il calcolo diretto della GPU non è semplice come scrivere codice per le CPU del dispositivo. Quando si lavora con il framework Metal di Apple, una serie approssimativa di operazioni per l’esecuzione del codice sulla GPU è simile a:

  • Trova un dispositivo GPU appropriato
  • Crea una coda per l’esecuzione dei comandi (ad esempio il MTLCommandQueue)
  • Avvolge i puntatori agli array di dati in un buffer strutturato; se i dati sono codice eseguibile, allora a stato del gasdottoaltrimenti a tampone regolare. Le GPU Apple utilizzano a spazio di memoria unificatoil che significa che non ne abbiamo bisogno in realtà copia tutti i dati nella memoria fisica specifica della GPU
  • Commettere il buffer dei comandi per l’esecuzione e attendere i risultati o impostare un gestore eventi al completamento
  • Estrarre i byte da un buffer di risposta e formattarli localmente con il codice di programma della CPU

La programmazione GPU grezza utilizza un modello asincrono.

Linguaggio di ombreggiatura dei metalli è un derivato di C++14 che può essere utilizzato per comporre logica personalizzata (chiamata “shader”) da eseguire su GPU compatibili con Metal. In generale, e se possibile, probabilmente staresti meglio usando il file quadro MPS (discusso più tardi) per funzionalità equivalenti quando possibile: tende ad essere altamente ottimizzato per le classi comuni di casi d’uso allineati alla GPU (come la moltiplicazione di matrici o reti neurali).

Il debug del codice MSL è piuttosto difficile. Puoi usare il Debugger di shader tramite Xcode, ma se desideri controllare o stampare valori intermedi senza Xcode, devi 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 attraverso il file kernel designazione. Il framework Metal passa gli ID per il contesto del thread chiamante corrente, o gruppo di thread, che può essere utilizzato per garantire scritture non sovrapposte. I thread possono essere rappresentati da un sistema ID tridimensionale; le dimensioni di questo spazio thread sono configurate nel codice wrapper.

Quello che segue è un implementazione della moltiplicazione di matrici naive algoritmo, combinato con alcuni srotolamenti del loop che sorprendentemente ne hanno migliorato significativamente le prestazioni. Questo è solo a scopo di confronto; normalmente il MPSMatrixMultiplication la funzionalità di MPS sarebbe più adatta.

kernel void matrix_multiply_naive(
device const MatrixParams *params,
constant float *A,
constant float *B,
device float *C,
// Indicates the thread's unique position within the entire grid of
// threads being executed. The uint2 type is a 2D coordinate, with
// fields x and y representing its indices on each axis.
// This parameter is not directly provided from the calling code,
// but provided by the Metal framework
uint2 gid ((thread_position_in_grid))
) {
if (gid.x >= params->a_rows || gid.y >= params->b_cols) {
return; // This thread is out of matrix dimensionality range, do nothing
}

float sum = 0.0;
int k;

// Loop unrolling; improves performance by a notable margin
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);
}

// Handle any remaining elements
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 a anche la funzione di trasposizione ingenua in MSL per il confronto. Data una matrice trasposta, questo è un banale adattamento alla logica di cui sopra, il cui ciclo interno corre sulle righe di B anziché lungo le sue colonne:

// Loop unrolling; improves performance by a notable margin
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); // Note this is gid.y * cols plus 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);
}

// Handle any remaining elements
for (; k < params->a_cols; ++k) {
sum += A(gid.x * params->a_cols + k) * B(gid.y * params->b_cols + k);
}

IO ho discusso questo approccio in un post precedente sul blog come un modo abbastanza semplice per migliorare le prestazioni scalari dell’algoritmo ingenuo, almeno sulle CPU. Ne parleremo più avanti.

La struttura Metal offre la possibilità di farlo compilare una libreria dal codice sorgente Metal. Una volta caricato il contenuto del file, il codice di collegamento cerca le funzioni del kernel per nomee inizializza un nuovo MTLComputePipelineState che rappresenta il codice funzione compilato.

id<MTLDevice> device = MTLCreateSystemDefaultDevice();

// Compile and initialize a new library located at the provided source path.
MTLCompileOptions *compileOptions = (MTLCompileOptions new);
compileOptions.languageVersion = MTLLanguageVersion3_0;

// Wrap input source path string
NSString *ss = (NSString stringWithUTF8String:source_path);

// Initialize new library containing compiled shader functions
id<MTLLibrary> lib = (device newLibraryWithSource:ss
options:compileOptions
error:&error);

// Create a representation of the naive multiplication public shader function in
// the Metal library created above
id<MTLFunction> naiveFunction =
(lib newFunctionWithName:@"matrix_multiply_naive");

// Create the new compute pipeline state
id<MTLComputePipelineState> pipelineStateNaive = (device newComputePipelineStateWithFunction:naiveFunction
error:&error);

Per chiamare effettivamente il codice Metal nativo, la configurazione del thread deve essere impostatoe i buffer della GPU sono inizializzati.

(computeEncoder setComputePipelineState:pipelineStateNaive);

MTLSize threadsPerGrid = MTLSizeMake(params->a_cols, params->a_rows, 1);

// Calculate a threadgroup size.
// https://developer.apple.com/documentation/metal/calculating_threadgroup_and_grid_sizes?language=objc
NSUInteger w = pipelineStateNaive.threadExecutionWidth;
NSUInteger h = pipelineStateNaive.maxTotalThreadsPerThreadgroup / w;
MTLSize threadsPerThreadgroup = MTLSizeMake(w, h, 1);

// Encode kernel function inputs
(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);

// Encode the compute command.
(computeEncoder dispatchThreads:threadsPerGrid
threadsPerThreadgroup:threadsPerThreadgroup);

// End the compute pass.
(computeEncoder endEncoding);

// Execute the command.
(commandBuffer commit);

Questo è molto, quindi illustrerò le relazioni qui:

Layout di alto livello di concetti, tipi e hardware all'interno del wrapper Objective-C
Layout di alto livello di concetti, tipi e hardware all’interno del wrapper Objective-C

IL Quadro MPS è una libreria ad alte prestazioni fornita da Apple per l’utilizzo con i suoi Famiglia di GPU Metal. Offre funzionalità dalle attività di immagine a supporto della rete neurale.

Le API sono disponibili principalmente tramite Swift o Objective-C, sebbene esista anche un file Metallo-cpp biblioteca disponibile per l’uso.

IL API di moltiplicazione MPSMatrix È ragionevolmente facile da usare. Come per il codice MSL sopra, i comandi MPS devono ancora essere codificati nel file MTLCommandBuffer e impegnato in modo asincrono per l’esecuzione.

// Define Matrix "descriptions", accounting for matrix dimensionality and byte size
MPSMatrixDescriptor *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);

// Output matrix
MPSMatrixDescriptor *descriptorC = (MPSMatrixDescriptor matrixDescriptorWithDimensions:a_rows
columns:b_cols
rowBytes:b_cols * sizeof(float)
dataType:MPSDataTypeFloat32);

// Initialize matrix representations using above descriptions and matrix buffers
MPSMatrix *matrixA = ((MPSMatrix alloc) initWithBuffer:bufferA descriptor:descriptorA);
MPSMatrix *matrixB = ((MPSMatrix alloc) initWithBuffer:bufferB descriptor:descriptorB);
MPSMatrix *matrixC = ((MPSMatrix alloc) initWithBuffer:bufferC descriptor:descriptorC);

// Creates the multiplication instance
MPSMatrixMultiplication *matrixMultiplication = ((MPSMatrixMultiplication alloc) initWithDevice:device
resultRows:a_rows
resultColumns:b_cols
interiorColumns:a_cols);

// Encodes the multiplication command into the command buffer for the GPU
id<MTLCommandBuffer> commandBuffer = (commandQueue commandBuffer);
(matrixMultiplication encodeToCommandBuffer:commandBuffer
leftMatrix:matrixA
rightMatrix:matrixB
resultMatrix:matrixC);

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

Vai è 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 di interfaccia di funzioni estranee.

La configurazione della direttiva è un po’ fragilema eventuali commenti immediatamente precedenti la riga import “C” (chiamato “il preamable”) verrà interpretato come importazione di intestazioni o argomenti di compilazione durante la compilazione del codice C di riferimento. Per 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 la riga di comando LDFLAGS
  • Compila il codice C con l’intestazione standard stdlib.h
  • Compila il codice C con l’intestazione del progetto locale metal.h

Sono stati necessari alcuni tentativi ed errori per far sì che il set corretto di flag del linker funzionasse su MacOS.

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

Si scopre che Apple racchiude un’implementazione BLAS nel suo Accelerate framework, quindi oltre a installare OpenBLAS tramite brewanche l’ubicazione della libreria deve essere fornita durante il collegamento:

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

IL go:embed La direttiva consente ai programmi Go di includere file in fase di compilazione, ovvero utile in questo caso quando vogliamo passare il contenuto del file sorgente MSL (mm.metal) al framework Metal, come discusso sopra, per la compilazione.

//go:embed mm.metal
var source string

// Compile the shader source code and initialize pipelines. The metalSource
// param contains the contents of an embedded Metal Shading Language file.
func Compile (metalSource string) {
// Wrap string in a C string
src := C.CString(metalSource)

// Free the above string after command queue is initialized
defer C.free(unsafe.Pointer(src))

// Compile the source, initialize pipelines and command queue
C.initializePipelineAndCommandQueue(src)
}

I riferimenti a C sopra si interfacciano con le API C tramite cgo, ad esempio:

// Calls initializeMTLBuffers from Obj-C bindings
C.initializeMTLBuffers(
a_data, // Input opaque pointer for A
b_data, // Input opaque pointer for B
C.int(4), // Converts 4 into C integer type
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),
}

// Return an unsafe pointer to this MatrixParams struct, cast to
// the native C representation defined in the shared header file
return (*C.MatrixParams)(unsafe.Pointer(&params));

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

Volevo confrontare le prestazioni della moltiplicazione di matrici basata su GPU con entrambe le implementazioni di livello superiore, come Biblioteca Gonumcosì come implementazioni intuitive, scritte a mano (e relativamente inefficienti).

Ho implementato un paio di algoritmi diversi in Go, incluso questo algoritmo ingenuo di trasposizione parallela, che divide ingenuamente il lavoro di moltiplicazione tra N goroutine:

func (a Matrix(T)) TransposeMultParallel(b *Matrix(T)) *Matrix(T) {
if a.Cols != b.Rows {
panic("matrices are the wrong size for multiplication")
}

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) // Add a count to the WaitGroup for the new goroutine
go func(i int) { // Kick off goroutine
defer wg.Done() // Decrease the count when the goroutine completes
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() // Wait for all goroutines to complete
return InitMatrixWithData(a.Rows, b.Cols, c_data)
}

Gonum BLAS è una libreria Go pura che implementa le interfacce BLAS. Tuttavia, può anche essere configurato per deviare le operazioni algebriche verso un’implementazione BLAS in codice nativo come OpenBLAS Attraverso netlib.

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

// Convert primitive arrays into gonum dense matrix types
gonum_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)

// Configure Gonum to use Gonum-default Go implementation
blas64.Use(gonum.Implementation{})

// Run a multiplication using Gonum BLAS impl
start = time.Now()
gonum_c.Mul(gonum_a, gonum_b)
bdata.TimeGonumNative(start)

// Configure Gonum to use Netlib which forwards operations to a
// native C-code BLAS implementation (OpenBLAS in our case)
blas64.Use(netlib.Implementation{})

// Run a multiplication using OpenBLAS impl through Gonum API
start = time.Now()
gonum_d.Mul(gonum_a, gonum_b)
bdata.TimeGonumOpenBLAS(start)

Mio codice di benchmarking esegue alcune prove per ciascuna delle seguenti implementazioni di moltiplicazione di matrici e riporta il tempo medio impiegato da ciascuna per moltiplicare due matrici quadrate di dimensionalità gradualmente crescente:

- Naive multiplication, in Go
- Transposed naive multiplication, in Go
- Goroutine-parallelized transposed naive multiplication, in Go
- Gonum pure Go-based BLAS multiplication
- Gonum-wrapped OpenBLAS multiplication, written in C
- Hand-implemented naive multiplication, in MSL, on GPU
- Hand-implemented transposed naive multiplication, in MSL, on GPU
- Metal Performance Shaders framework, called from Objective-C, on GPU

L’output del benchmarking è simile al seguente (i float sono ms):

2023-12-01 11:12:51.644 go-mm(75818:22427382) Using default device Apple M2
elements naive transpose transpose_parallel metal_naive 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
...

Alcuni complotti veloci Attraverso matplotlib

Grafico delle prestazioni di tutti gli approcci
Grafico delle prestazioni di tutti gli approcci

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

Visualizzazione della cronologia GPU di Activity Monitor: tutti gli approcci (l'asse Y è la percentuale di utilizzo)
Visualizzazione della cronologia GPU di Activity Monitor: tutti gli approcci (l’asse Y è la percentuale di utilizzo)

Potete vedere che la GPU non è particolarmente impegnata, perché il tempo viene dedicato principalmente alle operazioni della CPU. Ecco un’altra esecuzione, escluse le tre tecniche di moltiplicazione più lente:

Grafico delle prestazioni degli approcci escluse le mie varianti Go scritte a mano
Grafico delle prestazioni degli approcci escluse le mie varianti Go scritte a mano

Circa 16 milioni di elementi (4k x 4k), Gonum comincia a degradarsi. Puoi vedere chiaramente qui che il sistema basato su GPU e OpenBLAS le operazioni superano le implementazioni Go pure. Osservando solo gli approcci basati su GPU:

Grafico delle prestazioni delle operazioni di moltiplicazione di matrici eseguite solo sulla GPU
Grafico delle prestazioni delle operazioni di moltiplicazione di matrici eseguite solo sulla GPU

Un paio di note interessanti qui:

  • La libreria Metal Performance Shaders è sorprendentemente veloce
  • Non esiste una reale differenza di prestazioni tra l’approccio ingenuo e quello ingenuo trasposto

Per il secondo punto: questo è diverso dalle caratteristiche prestazionali della coppia di implementazioni basate su Go di cui sopra. Si scopre che i modelli di accesso alla cache favorevoli per le CPU non funzionano allo stesso modo per le GPU e il modo in cui i loro gruppi SIMD (o deformazioni) accedere alla memoria. Vedi l’utilizzo della GPU qui per il confronto:

Monitoraggio attività Visualizzazione della cronologia GPU: solo operazioni GPU
Monitoraggio attività Visualizzazione della cronologia GPU: solo operazioni GPU

Ora sto guardando solo OpenBLAS E MPS da solo: i due approcci più rapidi:

Grafico delle prestazioni di OpenBLAS rispetto all'API Metal Performance Shaders MPSMatrixMultiplication di Apple
Grafico delle prestazioni di OpenBLAS rispetto all’API Metal Performance Shaders MPSMatrixMultiplication di Apple

Circa 35 milioni di elementi, il OpenBLAS l’implementazione inizia a peggiorare, mentre MPS sta tenendo duro. La differenza qui è piuttosto notevole, con quest’ultimo che completa le stesse operazioni di moltiplicazione della matrice di 35 M elementi in <15% delle volte. È ragionevole supporre che la differenza continui a crescere con la cardinalità della matrice.

Naturalmente, ci sono probabilmente differenze algoritmiche tra questi due approcci, quindi questo non è un confronto equo tra CPU e GPU. Se tracciamo le differenze di prestazioni tra le mie due implementazioni codificate manualmente, appare così:

Grafico del rapporto prestazioni del mio codice di moltiplicazione della matrice scritto in MSL rispetto al mio codice scritto in Go
Grafico del rapporto prestazioni del mio codice di moltiplicazione della matrice scritto in MSL rispetto al mio codice scritto in Go

Ciò significa che l’ingenua implementazione basata su MSL completa la moltiplicazione di 5M elementi in appena l’1% delle volte della mia implementazione Go e tale rapporto sembra migliorare a favore della GPU nel tempo.

Fonte: towardsdatascience.com

Lascia un commento

Il tuo indirizzo email non sarà pubblicato. I campi obbligatori sono contrassegnati *