GPU acceleration of plane-wave density functional theory calculations in Abinit

Il documento presenta il porting su GPU del codice Abinit per calcoli DFT in onde piane, analizzando le prestazioni su architetture multi-GPU e confrontando l'efficienza di due algoritmi di diagonalizzazione, il metodo CG precondizionato a blocchi localmente ottimo e il filtraggio polinomiale di Chebyshev.

L'essenziale

Immagina di dover costruire un grattacielo di cristallo, ma invece di mattoni e cemento, devi usare milioni di onde invisibili (chiamate "onde elettroniche") per descrivere come si comportano gli atomi in un materiale. Questo è il compito del software Abinit, un programma scientifico che aiuta gli scienziati a progettare nuovi materiali, farmaci o batterie.

Fino a poco tempo fa, fare questi calcoli era come cercare di costruire quel grattacielo usando solo un vecchio computer da cucina: ci voleva una vita intera. Gli scienziati hanno quindi deciso di "trasferire" questo lavoro su GPU (le schede grafiche potenti che usiamo per i videogiochi), che sono come squadre di migliaia di operai specializzati pronti a lavorare tutti insieme.

Ecco cosa hanno fatto gli autori di questo articolo, spiegato in modo semplice:

1. Il Problema: Troppi Operai, Troppo Caos

I computer normali (CPU) sono come un capo cantiere molto intelligente che fa tutto da solo, passo dopo passo. Le GPU sono come un esercito di 10.000 operai che possono fare cose semplici molto velocemente, ma se il capo cantiere continua a dare ordini uno alla volta, gli operai restano fermi a guardare.

Il problema era che il codice di Abinit era scritto per il "capo cantiere" (CPU). Se lo avessimo semplicemente dato agli operai (GPU), si sarebbero bloccati perché dovevano aspettare istruzioni continue.

2. La Soluzione: Il "Metodo del Gruppo" (Batch Processing)

Per far funzionare le GPU, gli scienziati hanno dovuto cambiare strategia. Invece di chiedere agli operai di calcolare un'onda alla volta, hanno raggruppato le onde in grandi pacchi (chiamati "batch").

• Prima: "Tu, operaio 1, calcola l'onda A. Tu, operaio 2, calcola l'onda B..." (Molto lento, molta attesa).
• Ora: "Ehi, squadra! Ecco un pacco di 1000 onde. Calcolatele tutte insieme, tutte in una volta!"
  Questo permette alle GPU di lavorare al massimo della loro potenza, come una macchina da corsa che non deve mai frenare.

3. I Due Metodi per Risolvere il Puzzle

Per trovare la soluzione finale (la struttura stabile del materiale), il programma deve risolvere un'enorme equazione matematica. Hanno testato due metodi diversi per farlo:

• Metodo A (LOBPCG): È come cercare di indovinare la soluzione facendo piccoli passi, controllando ogni volta se ci si sta avvicinando, e correggendo la rotta. È preciso, ma richiede molti "controlli" (comunicazioni) tra gli operai, il che rallenta il lavoro sulle GPU.
• Metodo B (Filtraggio di Chebyshev): È come lanciare una rete da pesca gigante. Invece di controllare ogni singolo pesce, lanci la rete, la tiri su e poi selezioni solo i pesci che ti servono. Questo metodo fa un sacco di calcoli "bruti" (che le GPU amano) e fa pochissimi controlli.

La scoperta: Il Metodo B (Chebyshev) è stato il vincitore. Sfrutta la potenza delle GPU in modo molto più efficiente, come se avesse un motore turbo, mentre il Metodo A si inceppava un po' troppo spesso.

4. Il Risultato: Risparmio di Tempo ed Energia

Hanno fatto delle prove su supercomputer reali (con schede grafiche NVIDIA e AMD).

• Velocità: Usando le GPU, il lavoro è stato completato fino a 17 volte più velocemente rispetto ai computer normali.
• Energia: Non solo è più veloce, ma consuma meno energia per fare lo stesso lavoro. È come passare da un camion vecchio e rumoroso a un'auto elettrica silenziosa ed efficiente: arrivi prima e spendi meno benzina.

In Sintesi

Gli scienziati hanno preso un programma scientifico complesso, lo hanno "tradotto" per parlare la lingua delle schede grafiche moderne, e hanno scoperto che lavorare a gruppi grandi (batch) e usare il metodo della "rete da pesca" (Chebyshev) è la chiave per costruire il futuro dei materiali in tempi record.

Grazie a questo lavoro, in futuro potremo scoprire nuovi farmaci o batterie migliori molto più velocemente, risparmiando anche energia elettrica nel processo!

Sommario tecnico

Titolo: Accelerazione GPU dei calcoli di DFT a onde piane in Abinit

1. Il Problema

Le simulazioni su larga scala nella scienza dei materiali a livello quantistico, in particolare quelle basate sulla Teoria del Funzionale Densità (DFT) di Kohn-Sham, richiedono il calcolo di migliaia di stati elettronici. Questo comporta sfide significative nell'ottimizzazione delle architetture di High-Performance Computing (HPC) moderne, specialmente per ridurre i tempi di simulazione.
Il collo di bottiglia principale risiede nella risoluzione iterativa del problema agli autovalori (diagonalizzazione) per trovare le funzioni d'onda elettroniche. Queste operazioni coinvolgono:

• Algebra lineare su matrici di grandi dimensioni (Hamiltoniano).
• Trasformate di Fourier veloci (FFT).
• Operazioni di ortogonalizzazione e proiezione.
  I codici tradizionali, ottimizzati per CPU multi-core, non sfruttano appieno la potenza delle moderne architetture ibride CPU-GPU. Inoltre, la semplice portabilità del codice non è sufficiente; sono necessarie revisioni algoritmiche profonde per identificare operazioni matematiche efficienti su GPU e gestire la distribuzione della memoria delle funzioni d'onda.

2. Metodologia

Gli autori hanno riprogettato da zero la versione GPU del codice Abinit (una suite open-source per calcoli DFT a onde piane), adottando un approccio ibrido MPI + OpenMP con offloading su GPU.

Strategie Chiave:

• Batch Processing: Invece di elaborare le bande elettroniche singolarmente, i dati sono raggruppati in "batch" (lotti). Questo permette di eseguire operazioni su blocchi più grandi (es. moltiplicazioni matrice-matrice invece che vettore-matrice), massimizzando il parallelismo SIMD e riducendo l'overhead di memoria.
• Gestione della Memoria (GPU-Resident): La funzione d'onda viene trasferita dalla CPU alla GPU all'inizio di ogni iterazione del campo autoconsistente (SCF) e mantenuta interamente in memoria GPU durante la diagonalizzazione iterativa, minimizzando i trasferimenti Host-Device.
• Librerie Vendor: Utilizzo di librerie a basso livello fornite dai produttori (cuBLAS, cuSOLVER, cuFFT per NVIDIA; rocBLAS, rocSOLVER, rocFFT per AMD) tramite wrapper, evitando kernel CUDA/HIP personalizzati per garantire la portabilità.
• Livello di Astrazione: Implementazione di un livello di astrazione in Fortran che gestisce la complessità delle strutture dati distribuite, le comunicazioni MPI e l'interfaccia con le librerie GPU, rendendo il codice mantenibile ed estendibile.
• Algoritmi di Diagonalizzazione: Confronto e implementazione di due algoritmi iterativi:
  1. LOBPCG (Locally Optimal Block Preconditioned Conjugate Gradient): Basato su gradienti coniugati e ortogonalizzazione a blocchi.
  2. Filtraggio Polinomiale di Chebyshev: Basato su proprietà analitiche dei polinomi per filtrare lo spettro degli autovalori, evitando dipendenze tra blocchi durante l'iterazione.

3. Contributi Chiave

• Porting Completo: Presentazione di una strategia di porting che sposta il calcolo delle funzioni d'onda interamente sulla GPU, utilizzando modelli di programmazione moderni (OpenMP 5.0+).
• Analisi Teorica dei Costi: Sviluppo di modelli teorici per stimare il costo computazionale (FLOPS), le comunicazioni MPI e l'intensità aritmetica per LOBPCG e Chebyshev.
  • È stato dimostrato che il filtraggio di Chebyshev offre un'intensità aritmetica $k$ volte superiore (dove $k$ è il grado del polinomio) rispetto a LOBPCG, poiché applica l'operatore Hamiltoniano $k$ volte tra due comunicazioni MPI globali.
• Ottimizzazione Algoritmica: Dimostrazione che il filtraggio di Chebyshev è più adatto alle GPU rispetto a LOBPCG, poiché sposta il carico di lavoro verso operazioni "compute-bound" (applicate su grandi batch) e riduce la frequenza delle operazioni "memory-bound" e di comunicazione (come l'ortogonalizzazione e la procedura di Rayleigh-Ritz).
• Metriche di Prestazione: Introduzione di metriche specifiche per valutare l'efficienza energetica e le prestazioni su architetture eterogenee (NVIDIA e AMD).

4. Risultati

I benchmark sono stati eseguiti su supercomputer reali (Jean Zay con GPU NVIDIA, Adastra con GPU AMD, Topaze) su sistemi modello come un'interfaccia solido-liquido di Titanio (255 atomi, 4096 bande) e Ga2O3.

• Speedup e Scalabilità:
  • Le GPU NVIDIA (serie A100/H100) hanno mostrato speedup significativi rispetto alle CPU (fino a 17x su un singolo nodo per il passo di filtraggio).
  • Le GPU AMD (MI250X) hanno mostrato speedup inferiori, in parte a causa delle prestazioni della libreria ROCm nella procedura di Rayleigh-Ritz.
  • La scalabilità forte mostra che l'efficienza diminuisce con l'aumento dei nodi (legge di Amdahl), ma l'uso di GPU riduce il numero totale di nodi necessari rispetto a una configurazione solo CPU.
• Confronto Algoritmi (LOBPCG vs Chebyshev):
  • Il passo di Rayleigh-Ritz è il collo di bottiglia principale, essendo limitato dalla larghezza di banda della memoria (memory-bound) e dalle comunicazioni MPI.
  • Il filtraggio di Chebyshev permette di aumentare il grado del polinomio per migliorare la precisione senza degradare le prestazioni, sfruttando l'elevata intensità aritmetica delle applicazioni dell'Hamiltoniano su GPU.
  • Al contrario, aumentare le iterazioni in LOBPCG non porta benefici significativi e aumenta il carico computazionale inutile.
• Efficienza Energetica:
  • L'uso di GPU riduce drasticamente il consumo energetico totale rispetto ai cluster CPU-only (fattori di risparmio fino a 15x su NVIDIA).
  • Il consumo energetico su sistemi AMD scala quasi linearmente con il numero di nodi, mentre su NVIDIA rimane quasi costante, indicando una migliore efficienza di calcolo.
• Modello Roofline: L'analisi roofline conferma che le operazioni di filtraggio sono "compute-bound" (raggiungono la massima velocità di calcolo), mentre l'ortogonalizzazione e Rayleigh-Ritz sono "memory-bound".

5. Significato e Prospettive

Questo lavoro rappresenta un passo fondamentale verso l'uso efficiente delle architetture HPC eterogenee per la scienza dei materiali.

• Validazione: Dimostra che la semplice portabilità del codice non è sufficiente; è necessaria una riprogettazione algoritmica per sfruttare appieno le GPU.
• Scelta Algoritmica: Il paper stabilisce che per i calcoli DFT su larga scala su GPU, il filtraggio polinomiale di Chebyshev è superiore al metodo LOBPCG tradizionale, poiché minimizza le comunicazioni MPI e massimizza l'utilizzo delle unità di calcolo.
• Futuro: Gli autori suggeriscono l'adozione di metodi di "Spectrum Slicing" per ridurre ulteriormente la dimensione della procedura di Rayleigh-Ritz, permettendo la simulazione di sistemi ancora più grandi distribuendo le funzioni d'onda su più GPU in modo più efficiente.
• Disponibilità: La nuova implementazione è disponibile nella versione Abinit v10.6 e successive, rendendo queste tecnologie accessibili alla comunità scientifica.

In sintesi, il paper fornisce una guida tecnica completa su come trasformare un codice DFT legacy in una applicazione ad alte prestazioni su GPU, fornendo dati empirici e modelli teorici che guidano la scelta degli algoritmi per il futuro dell'HPC scientifico.