JZ-Tree: GPU friendly neighbour search and friends-of-friends with dual tree walks in JAX plus CUDA

Il paper introduce JZ-Tree, un'implementazione open-source basata su JAX e CUDA che utilizza una gerarchia di alberi "piano-based" in ordine Morton per abilitare una ricerca dei vicini e un clustering Friends-of-Friends ad alte prestazioni su GPU, superando di oltre un ordine di grandezza le librerie concorrenti per grandi dataset grazie a un accesso alla memoria coalesciato e a un'efficiente traversata duale degli alberi.

L'essenziale

Immagina di avere una biblioteca gigantesca con miliardi di libri (i dati) e di dover trovare, per ogni libro, gli altri 16 più simili a lui. Se fossi un bibliotecario umano che usa un metodo classico, dovresti controllare libro per libro, o forse usare un sistema di scaffali molto complesso (un "albero" di dati) che però, se devi farlo fare a un robot velocissimo (una GPU), si inceppa perché il robot deve saltare da uno scaffale all'altro in modo disordinato.

Questo è il problema che gli autori di questo articolo, Jens Stücker e il suo team, hanno risolto con il loro nuovo metodo chiamato JZ-TREE.

Ecco come funziona, spiegato con parole semplici e metafore:

1. Il Problema: Il "Salto" Disordinato

Le moderne schede grafiche (GPU) sono come squadre di migliaia di operai che lavorano tutti insieme. Sono velocissimi se tutti fanno la stessa cosa nello stesso momento (come spostare mattoni da un mucchio all'altro in fila).
I vecchi metodi per cercare i "vicini" (come gli alberi KD) costringono questi operai a saltare in direzioni diverse e a leggere i dati in modo disordinato. È come se in una catena di montaggio, ogni operaio dovesse correre in un magazzino diverso per prendere un pezzo: il traffico si blocca e la velocità crolla.

2. La Soluzione: La "Z-Order" (La Serpentina)

Gli autori hanno inventato un nuovo modo di organizzare i dati. Invece di costruire un albero con rami che si diramano in modo caotico, hanno deciso di ordinare tutti i punti (i libri, le stelle, le particelle) lungo una linea immaginaria a forma di "Z" (o serpentina).
Immagina di prendere una mappa della città e di tracciare un percorso a serpentina che tocca ogni casa una sola volta, senza mai saltare indietro.

• Vantaggio: Ora, quando i tuoi operai (la GPU) devono lavorare, possono camminare lungo questa linea a serpentina passo dopo passo. Non devono saltare da un lato all'altro della città. Tutto scorre fluido e ordinato.

3. L'Albero "Piatta" (I Piani e i Gruppi)

Invece di costruire un albero altissimo e sottile (dove devi fare molti salti per arrivare in fondo), hanno costruito una serie di "piani" o livelli.

• Immagina di avere un grattacielo. I metodi vecchi ti fanno salire e scendere le scale per trovare una stanza specifica.
• Il metodo JZ-TREE è come avere un edificio dove ogni piano è diviso in grandi stanze. Tu sai esattamente su quale piano sei e quanti passi devi fare. Non devi perdere tempo a cercare le scale.
• Come sono organizzate le stanze? Ogni stanza (o foglia dell'albero) contiene gruppi di al massimo 48 punti. La regola fondamentale è che se due punti sono vicini sulla linea a serpentina (Z-order), devono stare nella stessa stanza. Questo significa che le stanze possono essere piccole o grandi (fino a 48), ma non si spezzano mai i gruppi di punti che sono naturalmente vicini tra loro.
• Questo permette alla GPU di lavorare in "doppio": può controllare due gruppi di stanze contemporaneamente in modo perfetto, senza che gli operai si disturbino a vicenda.

4. Cosa fanno di concreto?

Hanno applicato questa idea a due compiti molto comuni nella scienza:

1. Cercare i vicini più prossimi (KNN): Come trovare i 16 amici più vicini a te in una folla di milioni di persone.
2. Gruppi di amici (FoF - Friends of Friends): Come raggruppare le persone in "clan" se sono abbastanza vicine tra loro (usato spesso per trovare galassie o ammassi di stelle nell'universo).

5. I Risultati: La Super Velocità

Grazie a questo metodo, il loro programma JZ-TREE è:

• 10 volte più veloce delle migliori tecnologie attuali quando si tratta di grandi quantità di dati (milioni o miliardi di punti).
• Scalabile: Funziona benissimo anche se usi 64 schede grafiche insieme, come se avessi un esercito di robot che lavora all'unisono senza mai andare in conflitto.

In Sintesi

Gli autori hanno detto: "Non costruiamo un labirinto complicato per i nostri robot. Costruiamo una strada dritta e ordinata (la Z-order) e diamo loro una mappa a piani dove i punti vicini restano sempre insieme in gruppi di massimo 48. In questo modo, la potenza mostruosa delle schede grafiche moderne viene finalmente sfruttata al 100%."

Hanno reso disponibile questo codice gratuitamente (JZ-TREE) per chiunque voglia fare simulazioni scientifiche veloci, come studiare l'universo o simulare il movimento di fluidi, senza più aspettare ore per ottenere i risultati.

Sommario tecnico

1. Il Problema

Gli algoritmi basati su alberi spaziali (come le KD-tree) sono lo standard per problemi di ricerca di vicini e clustering nell'HPC basato su CPU. Tuttavia, il trasferimento diretto di questi algoritmi alle architetture GPU spesso fallisce nel sfruttare la loro elevata capacità computazionale.
Le principali sfide sono:

• Divergenza dei thread: La natura ramificata della traversata degli alberi causa thread che seguono percorsi diversi, riducendo l'efficienza parallela.
• Accesso alla memoria irregolare: I pattern di accesso alla memoria non sono "coalesced" (contigui), il che è critico per le GPU dove le transazioni di memoria sono condivise tra gruppi di thread (warps).
• Costruzione costosa: La costruzione tradizionale degli alberi (top-down) è spesso troppo lenta o complessa da parallelizzare efficacemente su GPU.

2. Metodologia: JZ-TREE

Gli autori propongono JZ-TREE, un framework che implementa una gerarchia di alberi basata su ordinamento Morton (z-order) e progettata specificamente per le GPU.

A. Costruzione dell'Albero (Bottom-Up)

Invece di costruire un albero binario profondo e irregolare, JZ-TREE utilizza un approccio bottom-up:

1. Ordinamento Z-Order: I punti vengono ordinati lungo una curva di riempimento dello spazio (Morton). Per i dati in virgola mobile, viene utilizzato un operatore di confronto personalizzato basato sul bit più significativo differenziale (msb), evitando la necessità di discretizzare il dominio.
2. Gerarchia di "Tree-Plane": L'albero non è un albero binario classico, ma una serie di "piani" di nodi.
   • I nodi sono definiti come segmenti contigui nell'array ordinato.
   • La profondità dell'albero è fissa e piccola.
   • I nodi possono avere un numero variabile di figli e dimensioni irregolari, ma sono archiviati in modo contiguo in memoria.
3. Regolarizzazione e Costruzione delle Foglie: Per evitare nodi con estensione spaziale eccessiva (tipico in distribuzioni non uniformi), viene applicato un limite globale sulla dimensione del volume dei nodi.
   • Vincolo sulle foglie: Le foglie contengono al massimo 48 punti.
   • Vincolo di coerenza: Tutti i punti che risiedono nella stessa cella z-order devono essere mantenuti insieme nella stessa foglia. Questo garantisce che la struttura rispetti l'ordinamento spaziale senza spezzare gruppi di punti adiacenti nella curva di riempimento, anche se ciò comporta una variazione nelle dimensioni delle foglie (fino al limite di 48).

B. Traversata Dual-Tree

L'algoritmo utilizza una dual tree walk (traversata duale) ottimizzata:

• Collaborazione tra thread: I thread lavorano in gruppi per elaborare interazioni tra gruppi di nodi, massimizzando l'accesso coalesced alla memoria globale.
• Liste di interazione: Viene mantenuta una lista di interazioni tra nodi sorgente e nodi riceventi. L'algoritmo procede dal livello più grezzo (root) a quello più fine (foglie), filtrando le interazioni non necessarie basandosi su limiti di distanza inferiori (dlow) e superiori (dup).
• Implementazione JAX/CUDA: L'implementazione è scritta in JAX, sfruttando la compilazione JIT e le interfacce FFI (Foreign Function Interface) per lanciare kernel CUDA ottimizzati. Questo permette di combinare le prestazioni delle GPU con la differenziabilità automatica e la facilità di programmazione di JAX.

C. Supporto Multi-GPU

L'algoritmo è stato esteso per il calcolo distribuito:

• Ogni GPU mantiene i nodi locali e le relative liste di interazione.
• I dati dei nodi sorgente remoti necessari per le interazioni vengono richiesti in modo efficiente e sincronizzato tra i livelli dell'albero.
• La comunicazione è ridotta al minimo, avvenendo principalmente per la sincronizzazione dei dati necessari per i livelli superiori dell'albero.

3. Contributi Chiave

1. Nuova Gerarchia di Albero: Un approccio basato su piani (tree-planes) che garantisce profondità fissa e accesso contiguo alla memoria, risolvendo i problemi di divergenza e coalescenza tipici degli alberi su GPU. La struttura delle foglie è vincolata a un massimo di 48 punti, mantenendo la coerenza spaziale z-order.
2. Implementazione Open Source: Rilascio di JZ-TREE su GitHub e PyPI, che fornisce un framework riutilizzabile per algoritmi basati su alberi.
3. Algoritmi Specifici: Implementazione e ottimizzazione di due algoritmi fondamentali:
   • Ricerca esatta dei k-Nearest Neighbors (kNN).
   • Clustering Friends-of-Friends (FoF), cruciale per l'identificazione di aloni nella cosmologia.
4. Supporto Multi-Tipo: Capacità di gestire diversi tipi di punti (es. punti sorgente e punti di query distinti) all'interno di un'unica struttura ad albero condivisa.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su nodi NVIDIA A100 (cluster Leonardo di CINECA).

• Prestazioni kNN:
  • JZ-TREE supera di un ordine di grandezza (fattore >10) le librerie concorrenti più vicine (come CLOVER, FAISS, CuPy-KNN) per problemi di grandi dimensioni ($N \gtrsim 10^7$).
  • Rispetto alle librerie basate su CPU (SciPy KD-Tree), il miglioramento è di oltre due ordini di grandezza.
  • Le prestazioni sono robuste su diverse distribuzioni di dati (uniforme, Gaussiana, simulazioni cosmologiche reali).
• Scalabilità Multi-GPU:
  • L'algoritmo scala efficacemente fino a 64 GPU.
  • Per $N=10^8$ punti per GPU, il tempo di esecuzione rimane competitivo, con una perdita di efficienza di solo ~30% passando da 2 a 64 GPU.
• Clustering FoF:
  • Rispetto a GADGET4 (MPI) e JFOF (GPU), JZ-TREE-FOF è fino a 116 volte più veloce rispetto all'implementazione CPU di GADGET4 e 18 volte più veloce rispetto a JFOF per dataset di grandi dimensioni ($512^3$ particelle).
  • È in grado di calcolare cataloghi FoF per $2048^3$ punti su 64 GPU in circa 3 secondi.

5. Significato e Impatto

Il lavoro di JZ-TREE rappresenta un passo significativo verso l'adozione di algoritmi basati su alberi nell'HPC moderno basato su GPU.

• Efficienza Energetica e Temporale: Sfrutta appieno il throughput delle GPU moderne, riducendo drasticamente i tempi di calcolo per simulazioni scientifiche intensive (come la fluidodinamica particellare, la materia oscura auto-interagente e la ricerca di aloni cosmologici).
• Flessibilità: La struttura dell'albero non è vincolata a una partizione spaziale rigida, rendendola adattabile a vari tipi di dati e dimensioni, pur rispettando i vincoli di coerenza z-order.
• Integrazione Scientifica: L'uso di JAX facilita l'integrazione in pipeline di inferenza basata su simulazioni e permette l'uso della differenziabilità automatica, aprendo nuove strade per l'analisi dei dati scientifici.

In sintesi, JZ-TREE dimostra che è possibile progettare algoritmi di traversata di alberi che non solo sono compatibili con le GPU, ma che ne sfruttano l'architettura in modo superiore rispetto alle implementazioni CPU tradizionali, offrendo guadagni di prestazioni senza precedenti per problemi su larga scala.