Accelerating finite-element-based projector augmented-wave density functional theory calculations with scalable GPU-centric computational methods

Questo articolo presenta un metodo scalabile basato su GPU per proiettori ad onda aumentata agli elementi finiti (PAW-FE) che sfrutta innovazioni algoritmiche come l'aritmetica a precisione mista e l'iterazione di sottospazio filtrata con polinomi di Chebyshev per ottenere accelerazioni significative e prestazioni pronte per l'exascale in simulazioni di teoria del funzionale della densità su larga scala e di accuratezza chimica.

L'essenziale

Immagina di dover prevedere il comportamento di una macchina complessa, come un motore automobilistico o un nuovo tipo di batteria. Per farlo con precisione, è necessario comprendere il comportamento di ogni singolo elettrone all'interno dei materiali che costituiscono la macchina. Questo è il compito di un campo chiamato Teoria del Funzionale della Densità (DFT). È come cercare di simulare una pista da ballo massiccia e intricata, dove miliardi di elettroni si muovono all'unisono.

Per lungo tempo, gli scienziati hanno avuto un problema: simulare queste danze per piccoli gruppi di atomi è facile, ma non appena si tenta di simulare un sistema grande e complesso (come una minuscola nanoparticella metallica o un foglio di materiale attorcigliato), il computer viene sopraffatto. È come cercare di dirigere una danza per 100.000 persone utilizzando un metodo progettato per 100; le istruzioni si aggrovigliano, la memoria si riempie e la simulazione impiega un'eternità per completarsi.

Questo articolo introduce un nuovo metodo super veloce per eseguire queste simulazioni, specificamente progettato per i moderni computer potenti che utilizzano le GPU (gli stessi chip che alimentano i videogiochi di fascia alta e l'intelligenza artificiale). Ecco come hanno fatto, scomposto in concetti semplici:

1. Il Vecchio Metodo vs. La Nuova Mappa

• Il Vecchio Metodo (Onde Piane): Immagina di cercare di mappare una città utilizzando una griglia gigante e uniforme in cui ogni pollice quadrato ha le stesse dimensioni. Se vuoi vedere un dettaglio minuscolo (come un singolo mattone su un edificio), devi rendere l'intera griglia incredibilmente fine, anche per il cielo vuoto sopra la città. Questo spreca una quantità enorme di potenza di calcolo. È così che funziona la maggior parte del software attuale.
• Il Nuovo Metodo (Elementi Finiti): Gli autori utilizzano un approccio di "mappa intelligente". Immagina una mappa che ingrandisce solo dove è necessario (come il centro città affollato) e riduce la scala dove è vuoto (come il cielo). Questo è chiamato discretizzazione a Elementi Finiti (FE). Permette loro di concentrare la potenza di calcolo esattamente dove gli elettroni stanno facendo cose interessanti, risparmiando enormi quantità di tempo e memoria.

2. Il Trucco "PAW": Il Costume Magico

Per rendere la matematica ancora più semplice, utilizzano un metodo chiamato Onda Proiettata Arricchita (PAW).

• Il Problema: Gli elettroni vicino al centro di un atomo (il nucleo) si muovono e vibrano selvaggiamente, rendendoli difficili da calcolare.
• La Soluzione: Il PAW è come mettere un "costume liscio" sugli elettroni. Finge che gli elettroni siano lisci e facili da gestire per la maggior parte del calcolo, ma mantiene un "trucco magico" segreto che gli permette di rivelare istantaneamente il vero comportamento selvaggio degli elettroni proprio quando deve controllare i dettagli vicino al nucleo. Questo permette loro di utilizzare una mappa molto più grossolana (più semplice) senza perdere accuratezza.

3. L'Impulso di Velocità della GPU: La Catena di Montaggio

Gli autori non hanno solo cambiato la mappa; hanno modificato il modo in cui il computer esegue i calcoli per adattarsi alle moderne GPU.

• Il Collo di Bottiglia: Di solito, i computer trascorrono molto tempo in attesa che i dati si spostino tra la memoria e il processore.
• La Soluzione: Hanno ridisegnato la matematica in modo che il computer possa eseguire molti calcoli contemporaneamente (come una catena di montaggio) invece che uno alla volta. Hanno anche utilizzato una tecnica intelligente chiamata Filtraggio di Chebyshev, che è come un setaccio che separa rapidamente gli elettroni "importanti" da quelli "non importanti", così il computer non spreca tempo su quelli di cui non ha bisogno.

4. Le Scorciatoie "Abbastanza Buone" (Precisione Mista)

Questa è forse la parte più creativa.

• L'Analogia: Immagina di dipingere un enorme murale. Per lo sfondo del cielo, non hai bisogno di mescolare il colore con precisione microscopica; una miscela "abbastanza buona" va bene ed è molto più veloce. Hai bisogno di precisione estrema solo per i minuscoli dettagli di un viso.
• L'Applicazione: Gli autori hanno realizzato che per le parti del calcolo che hanno solo bisogno di ottenere la forma generale corretta, possono utilizzare matematica a precisione inferiore (come usare un righello con meno tacche). Questo è molto più veloce sui chip moderni. Passano solo alla matematica "ultra-precisa" per gli ultimi, critici passaggi.
• Il Risultato: Mescolando matematica ad alta precisione e a bassa precisione, e sovrapponendo i trasferimenti di dati ai calcoli (facendo due cose contemporaneamente), hanno fatto sì che la simulazione venisse eseguita da 8 a 20 volte più velocemente di prima.

5. Cosa Hanno Realizzato

L'articolo afferma che con questi nuovi metodi:

• Velocità: Ora possono simulare sistemi con 10.000 fino a 130.000 elettroni in un tempo pratico (da minuti a ore) sui supercomputer.
• Confronto: Il loro metodo è circa 8 volte più veloce del software standard leader (Quantum ESPRESSO) per sistemi di questa dimensione.
• Scala: Hanno eseguito con successo una simulazione di un materiale "bistrato attorcigliato" (due fogli di atomi attorcigliati insieme) contenente 130.000 elettroni. Questa è una dimensione che in precedenza era impossibile simulare con questo livello di accuratezza utilizzando metodi standard.

Riassunto

In breve, gli autori hanno costruito un nuovo motore altamente efficiente per la simulazione dei materiali. Hanno combinato una "mappa intelligente" che ingrandisce solo dove necessario, un trucco del "costume magico" per semplificare la matematica e una modalità "avanti veloce" che utilizza una precisione inferiore per i passaggi non critici. Il risultato è uno strumento in grado di modellare materiali massicci e complessi sui moderni supercomputer in una frazione del tempo che richiedeva in passato, aprendo la strada alla progettazione di nuovi materiali per batterie, elettronica e catalizzatori molto più rapidamente.

Sommario tecnico

1. Enunciato del Problema

La modellazione accurata di sistemi materiali complessi (ad esempio, interfacce, difetti, nanocluster, eterostrutture torsionali) richiede simulazioni di Teoria del Funzionale della Densità (DFT) che coinvolgono da $10^4$ a $10^5$ elettroni. Tuttavia, le implementazioni esistenti affrontano significativi colli di bottiglia:

• Limitazioni delle Onde Piane (PW): I codici PW-PAW ampiamente utilizzati (ad esempio, VASP, Quantum ESPRESSO) si basano su Trasformate di Fourier Veloci (FFT), che comportano comunicazioni massicce "tutti-con-tutti". Ciò li rende limitati dalla larghezza di banda e scarsamente scalabili sulle moderne architetture GPU exascale.
• Limitazioni dello Spazio Reale: I metodi nello spazio reale esistenti utilizzano spesso Pseudopotenziali a Conservazione del Normo (ONCV), che richiedono un gran numero di funzioni di base (alti gradi di libertà) per raggiungere l'accuratezza chimica, specialmente per i metalli di transizione, portando a ingenti footprint di memoria e costi computazionali.
• Incompatibilità Hardware: Gli attuali risolutori di autovalori sono spesso inadatti all'alta intensità aritmetica e ai requisiti di bassa latenza dei cluster GPU moderni, non riuscendo a sfruttare le capacità di precisione mista o a sovrapporre efficacemente la comunicazione al calcolo.

2. Metodologia

Gli autori presentano PAW-FE, una formulazione Projector Augmented-Wave (PAW) discretizzata agli elementi finiti (FE) ottimizzata per architetture GPU multi-nodo.

A. Formulazione Matematica

• PAW-GHEP: Il metodo risolve il Problema agli Autovalori Generalizzato Ermitiano (GHEP): $H\tilde{\Psi} = S\tilde{\Psi}\Lambda$, dove $H$ è l'Hamiltoniana e $S$ è la matrice di sovrapposizione PAW.
• Formalismo di Spin Collineare: Le equazioni sono derivate all'interno di un framework di spin collineare per gestire sistemi magnetici.
• Condizioni al Contorno: La discretizzazione FE accoglie naturalmente condizioni al contorno periodiche, semi-periodiche (strati) e non periodiche (nanocluster) senza artefatti di periodicità artificiale.

B. Innovazioni Computazionali

Per risolvere efficientemente il PAW-GHEP sulle GPU, gli autori hanno sviluppato diverse strategie algoritmiche chiave:

1. Iterazione del Sottospazio Filtrata di Chebyshev basata sul Residuo (R-ChFSI):

   • Invece del ChFSI standard, utilizzano una formulazione basata sul residuo ($R = H\tilde{\Psi} - S\tilde{\Psi}\Lambda$).
   • Vantaggio Chiave: Questa formulazione è tollerante ai prodotti matrice-vettore non esatti, consentendo l'uso di inversi approssimati per la matrice di sovrapposizione PAW ($S^{-1}$) e aritmetica a precisione ridotta senza sacrificare la convergenza.

2. Matrice di Sovrapposizione a Inverso Approssimato:

   • Invece di invertire esplicitamente la matrice sparsa $S$, utilizzano un'approssimazione computazionalmente economica ($\tilde{S}^{-1}$) derivata da un'approssimazione diagonale della matrice di massa e correzioni localizzate. Questo evita costose inversioni globali.

3. Quadratura Multi-Risoluzione:

   • Per gestire gli integrali centrati sull'atomo (densità pseudo-core) su griglie FE grossolane, viene impiegato uno schema di quadratura multi-risoluzione. Utilizza regole di quadratura raffinate solo all'interno delle sfere di augmentazione degli atomi, mentre utilizza regole più grossolane altrove, garantendo accuratezza senza raffinamento della mesh.

4. Precisione Mista e Comunicazione a Bassa Precisione:

   • Calcolo: Utilizza aritmetica FP32/TF32 per i passaggi di filtraggio di Chebyshev (moltiplicazioni matrice-matrice) e BF16 per la comunicazione tra vicini.
   • Robustezza: La dipendenza dell'algoritmo R-ChFSI dai residui garantisce che queste riduzioni di precisione non degradino l'accuratezza finale in doppia precisione dello stato fondamentale.

5. Sovrapposizione Calcolo-Comunicazione:

   • Il filtraggio di Chebyshev viene eseguito per blocchi. Mentre un blocco di funzioni d'onda subisce comunicazione inter-processore (MPI), un altro blocco esegue calcoli (operazioni GEMM) sulla GPU. Questo nasconde la latenza di comunicazione, un fattore critico per la scalabilità exascale.

6. Operazioni Dense a Livello di Cella:

   • Piuttosto che costruire matrici sparse globali, il metodo riformula le operazioni come moltiplicazioni dense matrice-matrice a livello di cella. Questo massimizza l'intensità aritmetica e la località della cache, rendendolo ideale per l'esecuzione GPU.

3. Contributi Chiave

• Formulazione PAW-FE: Una nuova formulazione PAW discretizzata agli elementi finiti all'interno di un formalismo di spin collineare che supporta condizioni al contorno generiche.
• Risolutore di Autovalori R-ChFSI: Un'estensione dell'Iterazione del Sottospazio Filtrata di Chebyshev basata sul Residuo per risolvere il PAW-GHEP, abilitando l'uso di inversi approssimati e precisione mista.
• Implementazione GPU Scalabile: Una strategia di implementazione completa che presenta aritmetica a precisione mista (FP32/TF32/BF16), sovrapposizione calcolo-comunicazione e algebra lineare densa a livello di cella.
• Integrazione Multi-Risoluzione: Uno schema di quadratura che permette una valutazione accurata degli integrali PAW centrati sull'atomo su mesh grossolane.

4. Risultati e Prestazioni

Il metodo è stato testato su supercomputer di classe leader (OLCF Frontier, ALCF Aurora, ALCF Polaris) utilizzando GPU AMD, Intel e NVIDIA.

• Accuratezza: Validato contro codici a onde piane (Abinit, Quantum ESPRESSO) per molecole ($O_2$, $NO_2$) e cristalli (BCC Cr). Gli errori in energia e forze sono ordini di grandezza al di sotto dei requisiti di accuratezza chimica ($O(10^{-12})$ Ha/atom per l'energia, $O(10^{-6})$ Ha/bohr per le forze).
• Accelerazione CPU-GPU: Ha raggiunto un'accelerazione da 8× a 20× sulle GPU rispetto alle esecuzioni solo CPU (architetture Intel e AMD).
• Confronto con Onde Piane (QE):
  • Per sistemi con ~10.000 elettroni, PAW-FE ottiene una riduzione di 8× nel tempo di parete minimo rispetto a Quantum ESPRESSO.
  • Per sistemi più grandi (>10.000 elettroni), l'accelerazione aumenta ulteriormente a causa della località delle funzioni di base FE rispetto alla comunicazione globale dei metodi PW.
• Confronto con ONCV-FE:
  • PAW-FE richiede ~6× meno risorse computazionali e ottiene un tempo di soluzione ~5× più veloce rispetto ai metodi FE che utilizzano Pseudopotenziali a Conservazione del Normo (ONCV), principalmente a causa dei gradi di libertà ridotti abilitati dal metodo PAW.
• Scalabilità Exascale:
  • Ha dimostrato con successo un calcolo dello stato fondamentale per un sistema di bistrato torsionale di WTe2 composto da 130.000 elettroni (11.000 atomi).
  • Ha raggiunto un tempo di soluzione di ~2 minuti per iterazione SCF su 400 nodi, dimostrando la fattibilità del metodo per simulazioni realistiche su larga scala.

5. Significato

Questo lavoro stabilisce PAW-FE come un metodo pronto per l'exascale per simulazioni di primo principio. Colmando il divario tra l'alta accuratezza del metodo PAW e l'efficienza degli elementi finiti nello spazio reale, supera i colli di bottiglia di comunicazione dei metodi a onde piane. Il dispiegamento riuscito di strategie a precisione mista e di sovrapposizione su diverse architetture GPU dimostra un percorso per simulare regolarmente sistemi materiali complessi con $10^5$ elettroni, abilitando nuove scoperte in catalisi, interfacce di batterie e materiali quantistici che erano precedentemente computazionalmente intrattabili.