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

Deze paper introduceert JZ-Tree, een open-source JAX-implementatie die een op Morton-codes gebaseerde 'plane-based tree hierarchy' gebruikt om efficiënte dual-tree traversals op GPU's mogelijk te maken, wat leidt tot een prestatieverbetering van meer dan een factor tien voor $k$-nabuurzoek en Friends-of-Friends-clustering vergeleken met bestaande GPU-bibliotheken.

De kern

Het Grote Probleem: De "Bomen" die niet in de computer passen

Stel je voor dat je een enorme bibliotheek hebt met miljarden boeken (dit zijn de data-punten). Je wilt voor elk boek snel de 16 dichtstbijzijnde boeken vinden. Of je wilt groepen boeken vinden die bij elkaar horen (zoals vrienden die bij elkaar wonen).

Op een normale computer (CPU) gebruiken wetenschappers al decennia een slimme methode: ze bouwen een boomstructuur. Ze verdelen de bibliotheek in grote secties, dan in kleinere, dan in nog kleinere, tot ze bij het juiste boek zijn. Dit werkt heel goed op een gewone computer.

Maar moderne supercomputers gebruiken GPU's (grafische kaarten). Een GPU is niet zoals een slimme, snelle professor die één ding heel goed doet; het is meer zoals een leger van 10.000 soldaten die allemaal tegelijk iets moeten doen.

• Het probleem: De traditionele boom-structuur is te "takkerig". Sommige soldaten moeten diep de boom in, anderen niet. Sommigen lopen naar links, anderen naar rechts. Omdat ze allemaal tegelijk moeten werken, raken ze in de war (ze "divergeren") en wachten ze op elkaar. Het is alsof je een leger probeert te laten marcheren door een doolhof waar iedereen een andere route neemt. Het resultaat: de GPU staat erbij en kijkt erbij, terwijl hij eigenlijk razendsnel zou moeten zijn.

De Oplossing: JZ-TREE (De "Vlakke" Boom)

De auteurs van dit paper (Jens Stücker en collega's) hebben een nieuwe manier bedacht om die boom te bouwen, speciaal voor die 10.000 soldaten. Ze noemen het JZ-TREE.

Hier is hoe het werkt, stap voor stap:

1. De "Z-Order" (De Slang)

In plaats van een boom te maken met takken die alle kanten op groeien, sorteren ze alle boeken eerst in één lange, rechte rij. Maar niet zomaar willekeurig. Ze gebruiken een trucje genaamd Morton-codering (of Z-order).

• De Analogie: Stel je voor dat je een doos met Lego-blokjes hebt. In plaats van ze in stapels te zetten, leg je ze in een slang die door de doos kronkelt. Als twee blokjes dicht bij elkaar liggen in de ruimte, liggen ze ook dicht bij elkaar in die slang.
• Door ze zo te sorteren, kunnen de soldaten (de GPU-threads) in één keer een stukje van de slang lezen. Ze hoeven niet meer te springen van links naar rechts. Dit heet "gecoalesceerde toegang" (alles in één keer oppakken).

2. De "Vlakke Boom" (Geen Diepe Gaten)

Traditionele bomen kunnen heel diep zijn. JZ-TREE bouwt geen diepe bomen, maar lagen (zoals verdiepingen in een flatgebouw).

• De Analogie: Denk aan een grote zaal met mensen.
  • Verdieping 0: De zaal is opgedeeld in kleine groepjes van maximaal 48 mensen.
  • Verdieping 1: Die groepjes worden samengevoegd tot grotere groepen van 384 mensen.
  • Verdieping 2: En zo verder, tot je heel grote groepen hebt.
• Belangrijk detail: Een groepje is niet altijd precies 48 mensen groot. Het kan kleiner zijn. De regel is: als mensen in de ruimte heel dicht bij elkaar zitten (in dezelfde "Z-Order cel"), moeten ze altijd samen in hetzelfde groepje blijven, ook al zijn er minder dan 48. De groepjes zijn dus maximaal 48 personen, maar gegarandeerd dat iedereen die bij elkaar hoort, ook echt bij elkaar zit.
• Het mooie is: elke verdieping heeft altijd evenveel lagen. De soldaten hoeven nooit te gissen hoeveel stappen ze moeten zetten. Ze weten precies waar ze zijn. Dit maakt het werk voor de GPU voorspelbaar en supersnel.

3. De Twee-Bomen Dans (Dual Tree Walk)

Nu moeten ze de zoektocht doen. Stel je hebt twee sets boeken: Set A (waar we zoeken) en Set B (waar we naar kijken).

• Oude methode: Je loopt door de boom van A en voor elk boek zoek je in de boom van B. Dit is traag.
• JZ-TREE methode: Ze laten de twee bomen "danssen". Ze kijken naar twee grote groepen (groepen op dezelfde verdieping).
  • Als twee groepen ver uit elkaar liggen, weten ze direct: "Geen noodzaak om te zoeken, deze groepen raken elkaar nooit." Ze gooien die groep weg.
  • Als ze dicht bij elkaar liggen, kijken ze naar de kleinere groepen eronder.
  • De kracht: Omdat de data zo goed geordend is, kunnen 32 soldaten tegelijk de relatie tussen twee grote groepen controleren. Het is alsof ze niet één voor één deuren openen, maar een hele muur tegelijk wegduwen.

Wat hebben ze bereikt?

De auteurs hebben dit getest op twee dingen:

1. Dichtstbijzijnde buren vinden (KNN): "Vind de 16 dichtstbijzijnde buren voor elke persoon in de stad."
2. Vrienden-van-vrienden (FoF): "Vind alle groepen mensen die binnen een bepaalde afstand van elkaar wonen."

De resultaten:

• Voor grote hoeveelheden data (miljoenen tot miljarden punten) is hun methode 10 tot 100 keer sneller dan de beste bestaande software.
• Ze kunnen dit probleem oplossen op één GPU, maar het werkt ook perfect als je 64 GPU's tegelijk gebruikt. Het schaalt bijna perfect: als je 64 keer meer kracht toevoegt, is het probleem 64 keer sneller opgelost.

Waarom is dit belangrijk?

Dit is niet alleen leuk voor wiskundige puzzels. Dit is cruciaal voor:

• Astronomie: Om te begrijpen hoe sterrenstelsels ontstaan en hoe donkere materie zich gedraagt in het heelal.
• Simulaties: Om duizenden keren dezelfde simulatie te draaien om te zien wat er gebeurt als je de instellingen een beetje verandert.
• AI en Machine Learning: Om snel patronen in enorme datasets te vinden.

Samenvatting in één zin

De auteurs hebben een manier bedacht om data te sorteren in een "slang" en in "vlakke lagen", zodat de duizenden kleine processors van een GPU niet hoeven te wachten op elkaar, maar als een goed geoliede machine razendsnel de dichtstbijzijnde buren of groepen kunnen vinden.

Het is alsof ze van een rommelige, chaotische bibliotheek een perfect georganiseerd magazijn hebben gemaakt waar robots alles in één seconde kunnen vinden.

Technische samenvatting

Titel en Context

Paper: JZ-TREE: GPU friendly neighbour search and friends-of-friends with dual tree walks in JAX/CUDA
Auteurs: Jens Stücker, Oliver Hahn, Lukas Winkler, Adrian Gutierrez Adame, en Thomas Flöss (Universiteit van Wenen).
Kernonderwerp: Een nieuwe, GPU-vriendelijke implementatie van ruimtelijke boom-algoritmen (specifically dual tree walks) voor nauwkeurige k-nabuurzoek (kNN) en Friends-of-Friends (FoF) clustering, geschreven in JAX met CUDA-kernen.

---

1. Het Probleem

Hoewel algoritmes gebaseerd op ruimtelijke boom-traversie (zoals KD-bomen) zeer efficiënt zijn op CPU's voor High-Performance Computing (HPC), presteren ze vaak slecht wanneer ze direct naar GPU-architecturen worden overgebracht. De redenen hiervoor zijn:

• Thread Divergentie: Boomalgoritmes hebben een vertakkende aard. Op GPU's, waar duizenden lichte threads parallel werken, leidt dit ertoe dat threads binnen een "warp" verschillende paden door de boom volgen. Dit vereist serialisatie en verlaagt de doorvoersnelheid drastisch.
• Irreguliere Geheugentoegang: Traditionele boomstructuren leiden tot niet-gecoalesceerde geheugentoegang (threads lezen niet-contigu geheugen), wat een primaire prestatie-beperking is op GPU's.
• Schaalbaarheid: Bestaande GPU-bibliotheken voor boomtraversie schalen vaak niet goed naar zeer grote datasets ($N \gtrsim 10^7$) of distribueerde multi-GPU-systemen.

---

2. Methodologie: JZ-TREE

De auteurs stellen een nieuw raamwerk voor dat specifiek is ontworpen voor GPU-architecturen, gebaseerd op een Morton (z-order) boomhiërarchie.

A. Boomconstructie (Bottom-Up)

In plaats van de gebruikelijke top-down benadering (zoals bij KD-bomen), gebruikt JZ-TREE een bottom-up constructie:

1. Z-order Sortering: De invoerpunten worden gesorteerd volgens een Morton-curve (z-order). Dit wordt gedaan via een aangepaste vergelijking die de meest significante verschillende bit van zwevende-kommagetallen analyseert, zonder precisie te verliezen.
2. Boom-Planes (Tree-Planes): De boom wordt niet opgebouwd als een diep geneste binaire boom, maar als een hiërarchie van "boom-planes".
   • Een plane bestaat uit knopen die de punten partitioneren.
   • De diepte van de boom is vast en klein, wat de traversie voorspelbaar maakt.
   • Knopen op een coarser plane kunnen een flexibel aantal kinderen hebben op de fijnere plane.
3. Bladknopen (Leaves) en Constraints:
   • De bladknopen bevatten maximaal 48 punten.
   • Er is een cruciale constraint: alle punten die binnen dezelfde z-order cel vallen, worden samen in dezelfde bladknoop gehouden.
   • Hierdoor variëren de bladgroottes (ze zijn niet altijd exact 48), maar zijn ze strikt begrensd door de bovengrens van 48. Dit voorkomt dat punten die ruimtelijk dicht bij elkaar liggen (volgens de z-order) over verschillende threads worden verdeeld, wat de efficiëntie ten goede komt.
4. Regulering: Om extreem grote knopen in gebieden met lage dichtheid te voorkomen, wordt een globale bovengrens voor het knoopvolume ingesteld.

B. Dual Tree Walk (Twee Boom-Traversie)

De kern van de methode is een dual tree walk die werkt op deze boom-planes:

• Interactie Lijsten: In plaats van individuele punten te vergelijken, worden interacties tussen groepen knopen (nodes) verwerkt.
• Coalesced Memory Access: Omdat de kinderen van een knoop contigu in het geheugen worden opgeslagen, kunnen threads binnen een warp gezamenlijk geheugen lezen. Dit maximaliseert de geheugenefficiëntie.
• Werkverdeling: Threads werken collaboratief om interacties tussen knopenparen te evalueren, wat redundante geheugentoegang minimaliseert.

C. Implementatie Details

• JAX & CUDA: De algoritmes zijn geïmplementeerd in JAX voor hoge-level programmering en JIT-compilatie, met zware berekeningen uitgevoerd in aangepaste CUDA-kernen via de FFI (Foreign Function Interface).
• Multi-GPU: Het systeem ondersteunt distributie over meerdere GPU's. Remote knopen worden slechts één keer opgevraagd per traversie-stap, wat communicatie-overhead minimaliseert.
• Meerdere Punttypes: Het systeem kan verschillende punttypes (bijv. bron- en query-punten) in één boom verwerken door hun posities te concateneren voordat er gesorteerd wordt.

---

3. Belangrijkste Bijdragen

1. Nieuwe Boomstructuur: Een platte, plane-gebaseerde boomhiërarchie die is geoptimaliseerd voor GPU-geheugentoegang en thread-parallelisme, in tegenstelling tot traditionele diepe bomen.
2. Hoogwaardige Implementaties: Effectieve implementaties van twee kritieke algoritmes:
   • Exacte k-Nearest Neighbour (kNN) Search.
   • Friends-of-Friends (FoF) Clustering (cruciaal voor kosmologische simulaties).
3. Open Source: De code is beschikbaar als JZ-TREE op GitHub en PyPI, wat een basis vormt voor toekomstige GPU-algoritmes.
4. Schaalbaarheid: Bewezen sterke schaling naar grote datasets en distribueerde multi-GPU-systemen.

---

4. Resultaten en Prestaties

De auteurs testen hun methode op het Leonardo-supercomputer (CINECA) met NVIDIA A100 GPU's.

• Prestatieverbetering: Voor grote datasets ($N \gtrsim 10^7$) behaalt JZ-TREE een prestatieverbetering van meer dan een orde van grootte (factor 10+) ten opzichte van de dichtstbijzijnde concurrerende GPU-bibliotheken (zoals FAISS, CLOVER, en GPU-KD-bomen).
  • Voorbeeld FoF: Bij $N=512^3$ punten is JZ-TREE ongeveer 5x sneller dan GADGET4 (32 CPU-kernen), 18x sneller dan JFOF (GPU), en 116x sneller dan GADGET4 (1 CPU-kern).
• Schaalbaarheid:
  • De methode schaal goed tot 64 GPU's.
  • Bij FoF-clustering op 64 GPU's kunnen $2048^3$ punten in ongeveer 3 seconden worden verwerkt.
  • De efficiëntie daalt slechts met een factor 2-3 bij schaling van 1 naar 64 GPU's, wat voornamelijk te wijten is aan communicatie-latentie bij het oplossen van globale links.
• Vergelijking met Bestaande Methoden:
  • CPU (KD-bomen): JZ-TREE is meer dan twee orders van grootte sneller dan CPU-gebaseerde KD-bomen (SciPy).
  • Anderen GPU-methoden: Bestaande GPU-KD-bomen en CLOVER (een graafgebaseerde methode) schalen kwadratisch of hebben hogere constante factoren, waardoor ze bij grote $N$ aanzienlijk trager worden.

---

5. Betekenis en Toekomstperspectief

• HPC Impact: De methode maakt het mogelijk om simulaties die duizenden herhalingen vereisen (zoals simulatie-gebaseerde inferentie) veel sneller uit te voeren, wat essentieel is voor moderne wetenschappelijke toepassingen.
• Flexibiliteit: Hoewel de focus ligt op lage dimensies ($d \approx 3$) en grote $N$, is het raamwerk uitbreidbaar naar andere boomgebaseerde algoritmes zoals DBSCAN, Fast Multipole Methods (FMM) en correlatiefunctie-schattingen.
• JAX-integratie: Door JAX te gebruiken, kunnen onderzoekers profiteren van automatische differentiatie en JIT-compilatie, wat de integratie in complexe wetenschappelijke workflows vergemakkelijkt.

Conclusie: JZ-TREE lost het fundamentele probleem op van het porten van boomalgoritmes naar GPU's door de data-layout en traversiestrategie volledig aan te passen aan de eisen van GPU-architecturen (coalesced memory, thread collaboration), wat leidt tot baanbrekende prestaties in ruimtelijke zoekopdrachten en clustering.