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

Die Arbeit stellt JZ-Tree vor, eine Open-Source-Implementierung in JAX und CUDA, die mithilfe einer für GPUs optimierten Morton-basierten Baumhierarchie und dualer Baumdurchläufe eine überlegene Leistung bei der Nachbarschaftssuche und Friends-of-Friends-Clustering im Vergleich zu bestehenden GPU-Bibliotheken erzielt.

Das Wesentliche

Stellen Sie sich vor, Sie sind der Chef einer riesigen Bibliothek mit Milliarden von Büchern (den Datenpunkten). Ihre Aufgabe ist es, für jedes Buch sofort die 16 nächsten Nachbarn zu finden – also die Bücher, die thematisch oder räumlich am ähnlichsten sind. Oder Sie müssen alle Bücher in Gruppen einteilen, die so nah beieinander stehen, dass sie sich fast berühren (das nennt man „Friends-of-Friends"-Clustering).

In der Welt der Computer-Wissenschaften ist das ein klassisches Problem. Normalerweise nutzen Supercomputer dafür eine Art „Baum-Struktur" (wie einen Stammbaum), um die Suche zu beschleunigen. Aber hier liegt das Problem: Computerchips für Grafiken (GPUs) sind wie ein riesiges Heer von 10.000 kleinen, schnellen Arbeitern, die alle gleichzeitig arbeiten sollen. Die alten Baum-Algorithmen wurden aber für einzelne, sehr kluge Arbeiter (CPUs) entwickelt. Wenn man diese alten Pläne einfach auf das Heer überträgt, entsteht Chaos: Die Arbeiter warten aufeinander, rennen durcheinander und greifen unordentlich auf die Regale zu. Das kostet enorm viel Zeit.

Die Autoren dieses Papers haben eine neue Lösung namens JZ-TREE erfunden. Hier ist die Erklärung, wie das funktioniert, mit ein paar einfachen Bildern:

1. Das Problem: Der unordentliche Bibliothekar

Stellen Sie sich vor, die alten Algorithmen wären wie ein Bibliothekar, der jedem Arbeiter eine völlig andere, zufällige Aufgabe gibt.

• Worker A muss zum Regal 1.
• Worker B muss zum Regal 1000.
• Worker C muss wieder zum Regal 1.

Da alle gleichzeitig arbeiten wollen, aber alle an unterschiedliche Orte rennen müssen, entsteht ein Stau. In der Computer-Sprache nennt man das „ineffiziente Speicherzugriffe". Die GPU wartet, statt zu rechnen.

2. Die Lösung: Der Zick-Zack-Pfad (Morton-Order)

Die Autoren haben einen genialen Trick angewendet: Sie haben alle Bücher (Datenpunkte) nicht nach alphabetischer Reihenfolge sortiert, sondern nach einer Zick-Zack-Linie, die sich durch den ganzen Raum windet (die sogenannte „Morton- oder Z-Ordnung").

• Die Analogie: Stellen Sie sich vor, Sie legen alle Bücher nicht in Reihen, sondern in einem einzigen, langen, gewundenen Gang aus, der sich wie eine Schlange durch die Bibliothek schlängelt.
• Der Vorteil: Wenn Sie nun eine Gruppe von Arbeitern schicken, um Bücher zu suchen, müssen sie alle hintereinander in diesem Gang entlanglaufen. Sie greifen auf Bücher zu, die direkt nebeneinander liegen. Das ist für die GPU wie ein gut geölter Förderband-Transport: Alles fließt reibungslos und schnell.

3. Die „Flache" Baum-Struktur

Statt einen tiefen, verzweigten Baum zu bauen (wie einen echten Baum mit vielen Ästen), bauen die Autoren eine Ebene für Ebene-Struktur.

• Die Analogie: Statt in einem mehrstöckigen Gebäude mit vielen Treppen und Wendeltreppen zu suchen (wo man ständig die Richtung ändern muss), bauen sie eine riesige, flache Halle mit klaren, geraden Gängen.
• Warum das hilft: Die GPU-Worker können in Gruppen (sogenannten „Warps") zusammenarbeiten. Sie schauen sich gemeinsam einen Abschnitt an, entscheiden, ob sie weitergehen müssen, und tun dies synchron. Kein Arbeiter muss warten, bis ein anderer eine Entscheidung trifft.

Ein entscheidendes Detail dabei ist, wie die Bücher in diesen Gängen gruppiert werden: Die Autoren bilden Gruppen von bis zu 48 Büchern. Wichtig ist jedoch: Diese Gruppen sind nicht starr auf genau 48 Bücher festgelegt. Wenn mehrere Bücher in einem bestimmten Abschnitt der Zick-Zack-Linie liegen, werden sie immer zusammen in eine Gruppe gepackt, auch wenn es weniger als 48 sind. Das garantiert, dass räumlich benachbarte Bücher niemals getrennt werden, während die Gruppengröße flexibel bleibt (maximal 48).

4. Die zwei Hauptaufgaben

Mit diesem neuen System haben die Autoren zwei Dinge extrem schnell gemacht:

• KNN (Die Suche nach den nächsten Nachbarn):

  • Szenario: „Wo sind die 16 nächsten Freunde von mir?"
  • Ergebnis: Das System findet diese Freunde für Milliarden von Punkten in Sekunden. Es ist 10- bis 100-mal schneller als die besten bisherigen Methoden auf Grafikkarten.
  • Warum? Weil die Arbeiter nicht mehr im Kreis rennen, sondern effizient entlang der Zick-Zack-Linie arbeiten.

• FoF (Friends-of-Friends Clustering):

  • Szenario: „Welche Sterne gehören zu derselben Galaxie?" (In der Kosmologie nutzt man das, um Galaxienhaufen zu finden).
  • Ergebnis: Auch hier ist das System extrem schnell. Es verbindet Punkte, die nah beieinander sind, zu Gruppen, ohne dabei die GPU zu überlasten.

5. Der große Vorteil: Skalierbarkeit

Das Beste an JZ-TREE ist, dass es nicht nur auf einem Computer funktioniert, sondern auf vielen Grafikkarten gleichzeitig (sogar auf 64 oder mehr).

• Die Analogie: Wenn die Bibliothek so groß wird, dass ein Team sie nicht mehr schafft, teilen sie die Zick-Zack-Linie einfach auf mehrere Teams auf. Da die Linie so vorhersehbar ist, wissen die Teams genau, wo sie aufhören und wo das nächste Team beginnt. Die Kommunikation zwischen den Teams ist minimal und effizient.

Zusammenfassung für den Alltag

Stellen Sie sich vor, Sie müssten eine riesige Party organisieren und jedem Gast die 16 nächsten Nachbarn zuweisen.

• Die alte Methode: Jeder Gast rennt wild durch den Raum, fragt zufällige Leute und wartet, bis jemand antwortet. Das dauert ewig.
• Die JZ-TREE-Methode: Alle Gäste stehen in einer langen, geordneten Schlange. Die Organisatoren gehen einfach die Schlange ab, schauen sich immer ein kleines Stück an und sagen: „Ihr drei hier, ihr seid Nachbarn." Das geht blitzschnell.

Das Fazit: Die Autoren haben einen Weg gefunden, wie man die rohe Rechenkraft moderner Grafikkarten (GPUs) wirklich ausnutzt, indem sie die Daten so anordnen, dass die Computer-Arbeiter nicht mehr stolpern, sondern im Takt marschieren. Das macht Simulationen in der Astrophysik, KI und anderen Bereichen viel schneller und energieeffizienter. Und das Beste: Der Code ist Open Source, also kann jeder ihn nutzen!

Technische Zusammenfassung

1. Problemstellung

Hochleistungsrechnen (HPC) verlagert sich zunehmend von CPUs auf GPUs, um von deren hoher arithmetischer Durchsatzleistung und Energieeffizienz zu profitieren. Viele etablierte Algorithmen für räumliche Suchen und Interaktionen, wie z. B. auf Bäumen basierende Methoden (z. B. KD-Bäume), sind auf CPUs hocheffizient. Die direkte Übertragung dieser Algorithmen auf GPU-Architekturen führt jedoch oft zu enttäuschenden Ergebnissen.

Die Hauptgründe für diese Ineffizienz sind:

• Thread-Divergenz: Die verzweigte Natur von Baum-Traversierungen führt dazu, dass Threads innerhalb eines Warps unterschiedliche Pfade nehmen, was die Ausführung serialisiert und die Leistung mindert.
• Irreguläre Speicherzugriffsmuster: Traditionelle Baumstrukturen führen zu nicht-kohärenten (uncoalesced) Speicherzugriffen, was den Speicherdurchsatz auf GPUs stark begrenzt.
• Hohe Konstruktionskosten: Das top-down Aufbauen von Bäumen ist auf GPUs oft teuer und schwer zu parallelisieren.

Ziel ist es, eine Baumstruktur zu entwickeln, die speziell für die Massivparallelität und den Speicherzugriff von GPUs optimiert ist, um exakte Nachbarschaftssuchen (k-NN) und Clustering-Algorithmen (Friends-of-Friends, FoF) effizient durchzuführen.

2. Methodik: JZ-TREE Framework

Die Autoren stellen JZ-TREE vor, ein Framework, das auf einer Morton-Order (z-Order) basierenden Baumhierarchie aufbaut. Der Kernansatz besteht darin, die Daten so zu strukturieren, dass sie für GPU-Kerne optimiert sind.

A. Baumkonstruktion (Bottom-Up)

Im Gegensatz zu klassischen top-down KD-Bäumen wird JZ-TREE in zwei Schritten konstruiert:

1. Morton-Sortierung: Die Eingabepunkte werden nach ihrer z-Order-Kurve sortiert. Dies geschieht durch Interleaving der Bits der Koordinaten (Morton-Encoding). Für Gleitkommazahlen wird ein benutzerdefinierter Vergleichsoperator verwendet, der die höchstwertigen unterschiedlichen Bits (MSB) der Mantissen und Exponenten analysiert, ohne die Präzision zu truncieren.
2. Bottom-Up-Verdichtung: Anstatt einen tief verschachtelten binären Baum zu erzeugen, wird eine Hierarchie von Baum-Ebenen (Tree-Planes) erstellt.
   • Die Ebene wird durch eine maximale Anzahl von Punkten pro Knoten ($N_{max}$) definiert.
   • Knoten werden so gebildet, dass sie zusammenhängende Segmente der sortierten z-Order-Liste umfassen.
   • Die Tiefe des Baumes ist überall fest und klein, was die Traversierung vorhersehbar macht.
   • Ein zentraler Aspekt ist die Behandlung der Blätter: Diese enthalten maximal 48 Punkte. Dabei gilt die zwingende Einschränkung, dass alle Punkte, die sich innerhalb derselben z-Order-Zelle befinden, zwingend im selben Blatt verbleiben müssen. Dies führt zu Blättern mit variabler Größe (bis zu 48 Punkten), anstatt einer starren Fixierung auf genau 48 Punkte pro Blatt.
   • Knoten können eine variable Anzahl von Kindern haben und ungleiche räumliche Ausdehnungen aufweisen.

B. Dual Tree Walk (Dualer Baum-Durchlauf)

Der Algorithmus nutzt einen dualen Baum-Durchlauf, um Interaktionen zwischen Gruppen von Knoten gemeinsam zu verarbeiten.

• Kollaborative Ausführung: Threads innerhalb eines Blocks arbeiten zusammen, um die Interaktionen zwischen Knotenpaaren zu bewerten.
• Speicherzugriff: Durch die flache, zusammenhängende Speicherlayout der Knoten (basierend auf der z-Order) werden Speicherzugriffe kohärent (coalesced) gestaltet. Dies maximiert den Speicherdurchsatz.
• Pruning: Es werden untere ($d_{low}$) und obere ($d_{up}$) Distanzschranken zwischen Knoten berechnet. Interaktionen, die sicher außerhalb des Suchradius liegen, werden frühzeitig verworfen.

C. Implementierungsdetails (JAX & CUDA)

• Die Rechenkerne sind in CUDA implementiert und werden über die Foreign Function Interface (FFI) von JAX aufgerufen.
• Dies ermöglicht Just-in-Time (JIT) Kompilierung, automatische Differentiation und eine hohe Abstraktionsebene.
• Für Multi-GPU-Szenarien wird ein Sampling-basierter Partitionierungsansatz verwendet, um die Last zu balancieren und Daten zwischen GPUs auszutauschen, wobei die Kommunikation minimiert wird.

3. Wichtige Beiträge

1. GPU-optimierte Baumhierarchie: Eine neue Baumstruktur, die auf z-Order-Sortierung und einer Ebene-basierten (plane-based) Hierarchie beruht, die Thread-Divergenz minimiert und Speicherzugriffe optimiert. Die Blattknoten folgen dabei der Regel, maximal 48 Punkte zu enthalten, wobei die Integrität von z-Order-Zellen gewahrt bleibt.
2. Implementierung von k-NN und FoF: Vollständige Implementierungen für exakte $k$-nächste-Nachbar-Suche und Friends-of-Friends-Clustering auf GPUs.
3. Open-Source-Software: Bereitstellung der Bibliothek JZ-TREE (verfügbar auf GitHub und PyPI), die als Basis für weitere baum-basierte Algorithmen (z. B. Fast Multipole Method, DBSCAN) dienen kann.
4. Skalierbarkeit: Demonstration einer starken Skalierung auf verteilte Multi-GPU-Systeme.

4. Ergebnisse

Die Autoren führten Benchmarks auf dem Leonardo-Supercomputer (CINECA) mit NVIDIA A100 GPUs durch.

• Leistungsgewinn: Für große Problemgrößen ($N \gtrsim 10^7$) erreicht JZ-TREE eine Verbesserung um mehr als eine Größenordnung (Order-of-Magnitude) im Vergleich zu den besten konkurrierenden GPU-Bibliotheken (wie FAISS, CLOVER, CuPy-KNN, JAXKD-CUDA).
• Vergleich mit CPU: Die CPU-basierte SciPy-Implementierung ist um mehr als zwei Größenordnungen langsamer.
• Skalierung:
  • Die Methode skaliert effizient von 1 bis 64 GPUs.
  • Bei der k-NN-Suche für $N=10^8$ Punkte auf 64 GPUs beträgt die Zeit ca. 1,3 Sekunden.
  • Beim FoF-Clustering (kosmologische Simulationen) ist JZ-TREE auf 64 GPUs etwa 5-mal schneller als GADGET4 (32 Kerne) und 116-mal schneller als GADGET4 (1 Kern).
• Robustheit: Die Leistung bleibt über verschiedene Verteilungen hinweg (uniform, normal, kosmologische Simulationen) stabil, wobei Periodizität (wrapping) nur einen geringen Overhead verursacht.

5. Bedeutung und Ausblick

Das Paper zeigt, dass baum-basierte Algorithmen, die traditionell als CPU-dominiert gelten, durch eine sorgfältige Anpassung der Datenlayout- und Traversierungsstrategie (insbesondere durch Nutzung der z-Order und kollaborativer Thread-Gruppen) auf GPUs extrem effizient laufen können.

• Anwendungsgebiete: Die Bibliothek ist direkt anwendbar in der Astrophysik und Kosmologie (z. B. Halo-Finding, N-Körper-Simulationen, Smoothed Particle Hydrodynamics) sowie in der maschinellen Lernforschung (Simulation-based Inference).
• Zukunft: Das Framework dient als Fundament für die Implementierung weiterer komplexer Algorithmen auf GPUs, wie z. B. der Fast Multipole Method (FMM), die in zukünftigen Arbeiten behandelt werden soll.

Zusammenfassend liefert JZ-TREE einen Weg, um die Lücke zwischen der theoretischen Komplexität von Baumalgorithmen und der praktischen Leistung auf modernen GPU-Hardware-Architekturen zu schließen.