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

Ce papier présente JZ-Tree, une bibliothèque open-source basée sur JAX et CUDA qui utilise une hiérarchie d'arbres Morton optimisée pour les GPU afin d'accélérer considérablement la recherche de voisins et le regroupement « amis d'amis » par rapport aux bibliothèques existantes.

L'essentiel

🌟 Le Problème : Des ordinateurs trop rapides pour leurs propres règles

Imaginez que vous avez un chef cuisinier ultra-rapide (un GPU, la puce graphique de votre ordinateur) capable de préparer des millions de plats en même temps. C'est le but des supercalculateurs modernes : faire des choses très vite.

Mais il y a un problème. La plupart des recettes (algorithmes) ont été écrites pour un chef qui travaille seul, lentement mais avec précision (un CPU). Si vous donnez une recette complexe et désordonnée à votre chef ultra-rapide, il va s'embrouiller. Il va passer son temps à attendre que les autres lui donnent des ingrédients, ou à changer de direction constamment. C'est ce qu'on appelle la "divergence" : tout le monde ne fait pas la même chose au même moment, et la vitesse s'effondre.

Les chercheurs de l'Université de Vienne ont dit : "Stop ! Il faut réécrire la recette pour que notre chef rapide puisse travailler sans s'arrêter."

🌳 La Solution : Une nouvelle façon de ranger les livres (JZ-TREE)

Pour trouver des voisins proches (comme dans une recherche de "qui habite le plus près de moi ?" ou pour regrouper des étoiles en galaxies), les ordinateurs utilisent souvent des arbres. Imaginez un arbre généalogique ou un classeur où l'on sépare les données en petits groupes.

Sur un ordinateur classique, cet arbre est comme une forêt dense et tortueuse. Pour un chef rapide qui veut travailler en équipe, c'est un cauchemar : certains doivent aller chercher un livre au fond du sous-sol, d'autres au grenier. Les allers-retours bloquent tout le monde.

JZ-TREE propose une nouvelle organisation, basée sur une astuce appelée l'ordre de Morton (ou courbe en "Z").

L'analogie du Supermarché

Imaginez un immense supermarché :

• L'ancienne méthode (CPU) : Les rayons sont organisés par catégorie (fruits, puis légumes, puis épicerie), mais à l'intérieur, les produits sont rangés au hasard. Si vous voulez trouver tous les produits d'un quartier précis, vous devez courir partout.
• La méthode JZ-TREE (GPU) : On réorganise tout le magasin en suivant une ligne en "Z" qui serpente de l'entrée jusqu'à la sortie. Les produits qui sont géographiquement proches les uns des autres dans le magasin sont aussi proches les uns des autres sur l'étagère.

Grâce à cette organisation, quand le chef rapide (le GPU) envoie ses 32 assistants (les "threads" d'un groupe) chercher des produits, ils peuvent tous avancer ensemble le long du rayon, sans jamais se croiser ni s'arrêter. C'est ce qu'on appelle un accès mémoire "coalescé" (tous ensemble, d'un seul coup).

Le secret de cette organisation réside dans la taille des paquets. Au lieu de créer des groupes de taille fixe, JZ-TREE crée des groupes d'au maximum 48 points. La règle d'or est simple : si plusieurs points se trouvent dans la même petite case de l'ordre "Z", ils doivent impérativement rester ensemble dans le même groupe, même si cela signifie que le groupe est plus petit que 48. Cela garantit que des points spatialement proches ne sont jamais séparés, tout en optimisant la taille des paquets pour la vitesse du GPU.

🚀 Les deux grandes applications

Les chercheurs ont testé cette nouvelle organisation sur deux tâches difficiles :

1. Trouver les voisins les plus proches (KNN) :

   • Le but : Pour chaque personne dans une foule, trouver les 10 personnes les plus proches.
   • Le résultat : Avec JZ-TREE, c'est 10 fois plus rapide que les meilleurs logiciels actuels quand on a des millions de personnes. C'est comme si on passait de 10 heures à 1 heure pour trier une foule immense.

2. Regrouper les amis (Friends-of-Friends) :

   • Le but : En cosmologie, on veut savoir quelles étoiles forment une galaxie. Si deux étoiles sont assez proches, elles sont "amies" et font partie du même groupe.
   • Le résultat : Là encore, le système est extrêmement rapide. Il permet de simuler l'univers entier en quelques secondes sur un seul ordinateur, alors que cela prenait des heures auparavant.

🤝 Travailler en équipe (Le "Dual Tree Walk")

Le secret de la vitesse, c'est comment les deux arbres (celui des "chercheurs" et celui des "cibles") se parlent.
Au lieu de demander à chaque chercheur de vérifier chaque cible un par un (ce qui serait lent), JZ-TREE permet aux équipes de dire : "Hé, notre groupe d'arbres est si loin de votre groupe d'arbres qu'on n'a même pas besoin de vérifier les détails !" ou "Notre groupe est si proche du vôtre qu'on peut tout vérifier ensemble !".

C'est comme si, au lieu de vérifier chaque maison d'un quartier pour voir si elle est proche d'une autre, on disait : "Tout le quartier A est à 5 km du quartier B, on peut ignorer les détails." Cela économise un temps fou.

🌍 Pourquoi c'est important pour nous ?

Ce n'est pas juste une victoire pour les mathématiciens.

• Pour l'Univers : Cela permet de simuler la formation des galaxies et de l'univers beaucoup plus vite, aidant les astronomes à comprendre d'où nous venons.
• Pour la médecine et l'IA : Ces mêmes techniques aident à analyser des images médicales ou à entraîner des intelligences artificielles plus efficacement.
• Pour l'avenir : Le code est libre et gratuit (open-source). N'importe quel développeur peut l'utiliser pour rendre ses propres applications plus rapides.

En résumé

Les chercheurs ont inventé une nouvelle façon de ranger les données (JZ-TREE) qui est parfaitement adaptée à la façon dont fonctionnent les puces graphiques modernes. Au lieu de forcer les puces rapides à travailler avec des vieilles méthodes lentes, ils ont créé une "autoroute" pour les données, en regroupant intelligemment les points proches (jusqu'à 48 par groupe) pour que le travail se fasse en flux continu.

Résultat ? Des calculs qui étaient autrefois impossibles ou très lents deviennent soudainement rapides comme l'éclair, permettant de résoudre des problèmes complexes de l'univers en un clin d'œil.

Résumé technique

1. Problématique

Les algorithmes basés sur la traversal d'arbres spatiaux (comme les arbres KD ou les octrees) sont la norme pour les problèmes de recherche de voisins et de clustering en calcul haute performance (HPC) sur CPU. Cependant, leur transposition directe sur les architectures GPU s'avère souvent contre-performante pour plusieurs raisons :

• Divergence des threads : La nature récursive et irrégulière de la traversal d'arbres crée une forte divergence de contrôle au sein des warps (groupes de threads), forçant une exécution sérialisée.
• Accès mémoire non coalescés : Les structures d'arbres classiques entraînent des motifs d'accès mémoire irréguliers, empêchant les transactions mémoire partagées entre threads, ce qui est crucial pour la performance GPU.
• Coût de construction : La construction d'arbres top-down est coûteuse et difficile à paralléliser efficacement sur GPU.

L'objectif est de concevoir une structure de données et des algorithmes adaptés spécifiquement aux contraintes des GPU (parallélisme massif, coalescence mémoire) pour des problèmes à grande échelle ($N \gtrsim 10^7$ points) en dimensions basses ($d \sim 3$).

2. Méthodologie

Les auteurs proposent JZ-TREE, un framework open-source implémenté en JAX et CUDA, reposant sur une hiérarchie d'arbres basée sur l'ordre de Morton (courbe Z).

A. Construction de l'arbre (Bottom-Up)

Au lieu d'une construction top-down, l'approche est bottom-up et repose sur deux étapes principales :

1. Tri Z-Order (Morton) : Les points sont triés selon une courbe de remplissage d'espace (Z-order). Pour les coordonnées flottantes, les auteurs utilisent un opérateur de comparaison personnalisé basé sur le bit significatif le plus élevé (MSB) qui diffère, évitant ainsi la troncature de précision requise par le codage entier classique.
2. Hiérarchie de "Tree-Planes" : Au lieu d'un arbre binaire profond, l'algorithme construit une série de "plans" (niveaux) de nœuds.
   • Les nœuds sont définis par des points de coupure (splitting points) dans l'index trié.
   • Chaque plan partitionne l'espace en cellules de Morton maximales contenant au plus 48 points.
   • Contrainte de cohérence : Tous les points appartenant à la même cellule Z-order doivent impérativement rester regroupés dans la même feuille. Par conséquent, la taille des feuilles varie selon la densité locale des données, mais est strictement bornée par 48 points.
   • La profondeur de l'arbre est fixe et faible, ce qui rend la traversal hautement prévisible.
   • Les enfants d'un nœud sont stockés de manière contiguë, permettant des lectures mémoire parfaitement coalescées.

B. Algorithmes de traversal Dual-Tree

L'approche utilise une traversal d'arbre dual (interaction entre deux arbres ou deux ensembles de points) optimisée pour le GPU :

• Listes d'interaction : Au lieu de traiter chaque paire de points, l'algorithme génère des listes d'interaction entre groupes de nœuds (leaf-to-leaf).
• Élagage collaboratif : Les threads d'un warp travaillent ensemble pour évaluer les interactions entre nœuds parents. Ils utilisent des bornes de distance ($d_{low}$ et $d_{up}$) pour élaguer les interactions inutiles.
• Gestion des mémoires : Les données des nœuds enfants sont chargées dans la mémoire partagée (shared memory) pour minimiser les accès à la mémoire globale. Des tas (heaps) de voisins sont maintenus dans les registres pour trouver les $k$ plus proches voisins.

C. Implémentation JAX/CUDA

• L'implémentation utilise des noyaux CUDA via l'interface FFI de JAX.
• Elle gère dynamiquement les allocations de mémoire (avec des estimations statiques pour le JIT) et supporte le multi-GPU via un échange de données de nœuds distants (remote nodes) uniquement lorsque nécessaire.

3. Contributions Clés

1. Structure de données "Tree-Planes" : Une hiérarchie d'arbres à profondeur fixe et à enfants contigus, conçue spécifiquement pour maximiser la coalescence mémoire et minimiser la divergence des threads sur GPU.
2. Algorithme de Tri Z-Order Flottant : Une méthode robuste pour trier des coordonnées flottantes sans perte de précision, essentielle pour une construction d'arbre fiable.
3. Implémentation Dual-Tree Optimisée : Des noyaux CUDA qui exploitent la collaboration des threads (via la mémoire partagée et les registres) pour traiter les interactions de groupes de nœuds, réduisant drastiquement les accès mémoire redondants.
4. Support Multi-GPU : Une architecture distribuée où chaque GPU gère ses nœuds locaux et demande les données des nœuds sources distants uniquement pour les niveaux de l'arbre pertinents, limitant la communication.
5. Bibliothèque Open-Source (JZ-TREE) : Une implémentation complète disponible sur GitHub et PyPI, intégrant la recherche de $k$-plus proches voisins (kNN) exacts et le clustering Friends-of-Friends (FoF).

4. Résultats Expérimentaux

Les benchmarks ont été réalisés sur des nœuds de calcul avec des GPU NVIDIA A100.

• Performance kNN (Recherche de voisins) :

  • Pour des problèmes de grande taille ($N \gtrsim 10^7$), JZ-TREE surpasse les bibliothèques concurrentes (FAISS, CLOVER, CUPY-KNN, SCIPY-KDTree) d'un ordre de grandeur (facteur > 10).
  • L'approche brute-force (FAISS) reste compétitive pour très petits $N$, mais JZ-TREE devient nettement supérieur dès que la complexité algorithmique prend le dessus.
  • L'efficacité de mise à l'échelle (scaling) est excellente, passant de 1 à 64 GPU avec une perte d'efficacité minime (environ 30% de perte entre 2 et 64 GPU).

• Performance FoF (Clustering) :

  • Sur des distributions cosmologiques réalistes, JZ-TREE est 5 à 100 fois plus rapide que les implémentations CPU (GADGET4, HFOF) et l'implémentation GPU existante (JFOF).
  • Exemple : Pour $512^3$ particules, JZ-TREE prend ~1.2s contre ~22s pour JFOF et ~82s pour HFOF (CPU).

• Robustesse : Les performances restent stables sur différentes distributions de données (grilles uniformes, distributions gaussiennes, simulations cosmologiques avec conditions aux limites périodiques).

5. Signification et Impact

Ce travail démontre que les algorithmes basés sur les arbres, traditionnellement réservés aux CPU, peuvent être réinventés pour exploiter pleinement la puissance des GPU.

• Changement de paradigme : Il prouve qu'il n'est pas nécessaire d'abandonner les arbres pour des méthodes "brute-force" ou approximatives sur GPU. Une conception adaptée (mémoire coalescée, profondeur fixe) permet d'obtenir des performances exactes et supérieures.
• Applications HPC : La bibliothèque JZ-TREE est directement applicable à des domaines exigeants comme l'hydrodynamique des particules lissées (SPH), la matière noire auto-interagissante et la recherche d'halos dans les simulations cosmologiques.
• Fondation pour l'avenir : Le framework JZ-TREE sert de base pour d'autres algorithmes complexes basés sur les arbres, tels que la méthode multipolaire rapide (FMM), dont la publication est annoncée.

En résumé, JZ-TREE comble le fossé de performance entre les algorithmes d'arbres spatiaux et l'architecture GPU, offrant une solution scalable, exacte et open-source pour le traitement de données massives en dimensions basses.