LP-GEMM: Integrating Layout Propagation into GEMM Operations

1. Il Problema

Nel calcolo scientifico e nei moderni carichi di lavoro di Machine Learning (ML), come le reti neurali profonde (MLP e Transformer), le operazioni di moltiplicazione di matrici (GEMM - General Matrix Multiply) sono spesso eseguite in sequenze dipendenti. Le librerie BLAS (Basic Linear Algebra Subprograms) allo stato dell'arte, come OpenBLAS, ottimizzano aggressivamente ogni singola chiamata GEMM. Tuttavia, sono vincolate dall'API BLAS, che richiede che ogni chiamata:

1. Riceva matrici di input in un layout di memoria canonico (es. row-major o column-major).
2. Restituisca l'output nello stesso layout canonico.

Quando si esegue una sequenza di GEMM (dove l'output di uno è l'input del successivo), questo vincolo impone che ogni operazione debba impacchettare (pack) i dati in un formato ottimizzato per il calcolo interno, eseguirlo, e poi disimballare (unpack) i risultati nel layout canonico prima che la successiva operazione possa iniziare. Questo ciclo continuo di packing e unpacking ridondante spreca risorse computazionali e di memoria, creando un collo di bottiglia significativo, specialmente in carichi di lavoro con catene lunghe di operazioni dipendenti.

2. Metodologia: LP-GEMM

Gli autori propongono LP-GEMM (Layout Propagation GEMM), una decomposizione del kernel GEMM di OpenBLAS progettata per propagare il layout dei dati tra operazioni sequenziali, eliminando la necessità di riscrivere i dati in formato canonico tra un'operazione e l'altra.

L'approccio si basa sulla suddivisione della sequenza di GEMM in tre tipi di kernel specializzati:

1. Initial GEMM (ini-GEMM): Esegue il primo passo di impacchettamento dei dati di input nel layout interno ottimizzato e propaga questo layout all'output.
2. Intermediate GEMM (mid-GEMM): Riceve i dati già impacchettati dall'operazione precedente. Non esegue alcun nuovo packing o unpacking; mantiene il layout interno e propaga l'output nello stesso formato per l'operazione successiva.
3. Ending GEMM (end-GEMM): È l'ultima operazione della catena. Riceve i dati nel layout propagato e li disimballa (unpack) per restituire il risultato nel layout canonico richiesto dall'utente o dalla successiva fase non-GEMM.

Aspetti Tecnici Chiave:

• Unificazione dei Layout: LP-GEMM allinea il layout di lettura (packed), l'ordine di calcolo (iterazione) e il layout di scrittura finale. Questo permette di evitare i passaggi di riordinamento dei dati.
• Micro-kernel: Vengono utilizzati due micro-kernel: uno "Propagate-Layout" che mantiene il formato interno e uno "Default" che ripristina il layout BLAS standard.
• Adattabilità: Il sistema gestisce anche carichi di lavoro complessi come i layer di attenzione dei Transformer, adattando operazioni intermedie (come Softmax e RoPE - Rotatory Position Embedding) per funzionare efficientemente con il layout propagato, minimizzando gli accessi alla memoria irregolari.
• Indipendenza dall'architettura: L'approccio è stato implementato e testato su architetture x86 (con estensioni AVX-512) e RISC-V (con estensioni RVV 1.0).

3. Contributi Chiave

• Decomposizione del Kernel: Introduzione di una nuova struttura di kernel (ini, mid, end) che sfrutta la propagazione del layout dei dati, rompendo il paradigma isolato delle chiamate BLAS tradizionali.
• Implementazione Pratica: Sviluppo di una versione C++ standalone del percorso di inferenza di Llama-3.2 che utilizza esclusivamente chiamate GEMM a livello BLAS, dimostrando la fattibilità dell'approccio in scenari reali.
• Ottimizzazione di Operazioni ML: Implementazione di operazioni comuni nei carichi di lavoro ML (Softmax, RoPE, RMSNorm) ottimizzate per il layout propagato su architetture AVX-512.
• Valutazione Multi-Architettura: Analisi approfondita delle prestazioni su piattaforme x86 e RISC-V, fornendo dati comparativi con librerie proprietarie e open-source.

4. Risultati Sperimentali

Le valutazioni sono state condotte su un processore Intel Xeon Gold 6252 (x86) e su un processore SpacemiT X60 (RISC-V).

• Prestazioni su GEMM Sequenziali:
  • Su x86, LP-GEMM ha ottenuto un speedup medio di 2.25x rispetto a OpenBLAS per sequenze di GEMM.
  • Su RISC-V, i guadagni sono stati ancora più significativi, con speedup fino a 5x per carichi di lavoro simili all'attenzione (Attention-like), a causa dell'inefficienza del kernel di unpacking di OpenBLAS su questa architettura.
  • Le prestazioni sono competitive rispetto a librerie ottimizzate dal vendor come Intel MKL, specialmente quando si considera l'intera catena di operazioni.
• Analisi del Layer di Attenzione (LLaMA 3.2):
  • L'integrazione di LP-GEMM nel layer di attenzione ha mostrato speedup fino a 2x su x86 per sequenze di input più piccole.
  • Su RISC-V, il miglioramento è cresciuto linearmente con la dimensione della sequenza (ntokens), poiché LP-GEMM elimina completamente il costo degli accessi alla memoria disordinati tipici dell'unpacking di OpenBLAS.
• Confronto con FlashGEMM: In benchmark di tre GEMM consecutivi, LP-GEMM ha superato sia OpenBLAS che FlashGEMM, dimostrando che la propagazione del layout è efficace anche senza fusione aggressiva di kernel complessi.

5. Significato e Implicazioni

Il lavoro di LP-GEMM segna un cambio di paradigma nell'ottimizzazione delle librerie lineari algebriche:

• Superamento dei Limiti BLAS: Dimostra che i vincoli dell'API BLAS, sebbene utili per la portabilità, possono essere aggirati in contesti di calcolo specifico per ottenere guadagni di prestazioni sostanziali senza sacrificare la correttezza semantica ai bordi del sistema.
• Efficienza Energetica e di Memoria: Eliminando il movimento ridondante dei dati (packing/unpacking), si riduce il traffico sulla memoria e il consumo energetico, fattori critici per l'Edge AI e i dispositivi con risorse limitate.
• Portabilità: L'approccio non richiede un tuning microarchitettonico profondo specifico per ogni nuova CPU, ma si basa su una strategia logica di gestione dei dati che funziona bene su diverse architetture vettoriali (x86 e RISC-V).
• Futuro della Ricerca: Suggerisce una direzione futura verso l'ottimizzazione "end-to-end" dei carichi di lavoro, estendendo la propagazione del layout anche ad altre operazioni oltre il GEMM, come le convoluzioni, integrando queste ottimizzazioni direttamente nei framework di compilazione.

In sintesi, LP-GEMM dimostra che comprendere l'interazione tra le operazioni all'interno di una pipeline di calcolo è fondamentale per sbloccare efficienze che le ottimizzazioni dei singoli kernel non possono raggiungere.