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

El artículo presenta JZ-Tree, una implementación de código abierto en JAX y CUDA que utiliza una jerarquía de árbol basada en planos de orden Morton para superar los problemas de divergencia de hilos y acceso a memoria en GPUs, logrando mejoras de rendimiento de más de un orden de magnitud en búsquedas de vecinos más cercanos y agrupamiento "friends-of-friends" para grandes conjuntos de datos.

Lo esencial

¡Claro que sí! Imagina que tienes una biblioteca inmensa con millones de libros desordenados y necesitas encontrar rápidamente los 10 libros más cercanos a uno que tienes en la mano, o agrupar todos los libros que están a menos de un metro de distancia entre sí.

Hacer esto en una computadora normal (CPU) es como tener un bibliotecario muy inteligente que revisa los libros uno por uno, pero en una computadora moderna con tarjetas gráficas (GPU), tenemos miles de bibliotecarios trabajando al mismo tiempo. El problema es que si les das las instrucciones de la biblioteca tradicional (árboles de búsqueda), se vuelven locos: algunos corren a la sección de ficción, otros a la de historia, y se chocan entre sí, perdiendo tiempo.

Aquí es donde entra el JZ-TREE, el nuevo método que presentan los autores de este artículo. Vamos a desglosarlo con analogías sencillas:

1. El Problema: El Caos en la Biblioteca

Las tarjetas gráficas (GPUs) son máquinas increíbles, diseñadas para hacer millones de tareas simples al mismo tiempo (como pintar millones de píxeles en un videojuego). Pero los algoritmos antiguos para buscar vecinos (como los "árboles KD") están hechos para pensadores lentos y metódicos (las CPUs).

Cuando intentas usar esos algoritmos viejos en una GPU, ocurren dos cosas malas:

• Confusión (Divergencia): Un grupo de bibliotecarios recibe una instrucción, pero algunos deben ir a la izquierda y otros a la derecha. La GPU tiene que esperar a que todos terminen su camino antes de seguir, desperdiciando potencia.
• Caminos tortuosos: Los datos no están ordenados en la memoria de forma que los bibliotecarios puedan agarrarlos todos de una sola vez. Tienen que correr de un lado a otro, perdiendo tiempo.

2. La Solución: La "Línea de la Serpiente" (Orden Z)

Los autores proponen una forma nueva de organizar los datos llamada Orden Z (o curva de Morton).

• La analogía: Imagina que en lugar de poner los libros en estantes rectos (1, 2, 3, 4...), los organizas en una larga cinta transportadora que hace zig-zag por toda la biblioteca.
• El truco: Si dos libros están físicamente cerca en la biblioteca, también estarán muy cerca en la cinta transportadora.
• El resultado: Ahora, cuando le das una orden a un grupo de 32 bibliotecarios (un "warp" en lenguaje técnico), todos pueden agarrar sus libros de la cinta transportadora al mismo tiempo, sin chocar y sin correr de un lado a otro. ¡Es como si todos tomaran un sorbo de agua de la misma fuente al mismo tiempo!

3. La Estructura: Pisos en lugar de Árboles Profundos

En lugar de construir un árbol gigante y profundo donde tienes que bajar escalón por escalón (lo cual es lento en paralelo), construyen "pisos" o niveles planos con una organización muy específica:

• La analogía: Imagina un edificio de apartamentos. En lugar de que cada piso tenga una cantidad diferente de habitaciones y escaleras locas, organizan el edificio en pisos donde los libros se agrupan en paquetes manejables.
• La regla de oro: Cada grupo (o "hoja" del árbol) contiene hasta 48 libros, pero con una condición importante: si varios libros caen en la misma sección de la cinta transportadora (misma celda de orden Z), todos deben ir juntos en el mismo grupo, sin importar cuántos sean.
• La ventaja: Esto significa que los grupos tienen un tamaño máximo de 48, pero pueden ser más pequeños si es necesario para mantener a los libros vecinos juntos. Todos los bibliotecarios saben exactamente dónde mirar y trabajan en equipo de forma muy eficiente, revisando grupos de libros juntos en lugar de uno por uno.

4. Los Dos Grandes Logros

Con esta nueva organización, probaron dos tareas difíciles:

• A. Encontrar los Vecinos Más Cercanos (KNN):

  • El reto: "¿Quiénes son los 16 libros más cercanos a este?"
  • El resultado: El JZ-TREE es más de 10 veces más rápido que las mejores herramientas actuales cuando hay millones de libros. Es como si antes tardaras 10 minutos en encontrar a tus amigos en un concierto y ahora tardaras 1 minuto.

• B. Agrupar a los Amigos (Friends-of-Friends):

  • El reto: "¿Qué libros forman grupos si están a menos de 1 metro de distancia?" (Esto es muy usado en astronomía para encontrar cúmulos de galaxias).
  • El resultado: De nuevo, una velocidad increíble. Pueden procesar miles de millones de puntos en segundos, algo que antes tomaba minutos o incluso horas en computadoras normales.

5. ¿Por qué es importante esto?

Este trabajo es como inventar un nuevo sistema de tráfico para una ciudad futurista llena de coches autónomos (las GPUs). Antes, el tráfico estaba diseñado para coches de la vieja escuela (CPUs), lo que causaba atascos.

• Para la ciencia: Permite a los astrónomos simular el universo entero en segundos en lugar de días.
• Para la inteligencia artificial: Hace que entrenar modelos sea más rápido y eficiente.
• Para todos: Es un código abierto (gratuito), lo que significa que cualquiera puede usarlo para hacer sus propias simulaciones más rápidas.

En resumen:
Los autores tomaron un problema difícil (buscar cosas en un espacio gigante) y lo resolvieron reorganizando la "biblioteca" de datos de una manera que las tarjetas gráficas modernas aman. En lugar de forzar a la GPU a pensar como una CPU lenta, le dieron instrucciones que se adaptan a su naturaleza masiva y paralela, logrando velocidades que antes parecían imposibles. ¡Es como pasar de caminar por la ciudad a volar en un cohete! 🚀

Resumen técnico

Aquí presento un resumen técnico detallado del artículo "JZ-TREE: GPU friendly neighbour search and friends-of-friends with dual tree walks in JAX plus CUDA", estructurado según los puntos solicitados y corregido según la retroalimentación de los autores.

1. El Problema

Los algoritmos basados en árboles espaciales (como KD-trees u Octrees) son el estándar en la computación de alto rendimiento (HPC) basada en CPU para problemas de búsqueda espacial y agrupamiento. Sin embargo, transferir directamente estos algoritmos a arquitecturas de GPU suele resultar en ganancias de rendimiento mucho menores de las esperadas debido a:

• Divergencia de hilos: La naturaleza ramificada de la traversión de árboles hace que los hilos dentro de un warp sigan caminos diferentes, obligando a la ejecución serializada y reduciendo la eficiencia.
• Patrones de acceso a memoria irregulares: Las estructuras de árbol tradicionales provocan accesos a memoria no coalescidos, lo que es crítico en GPUs donde el rendimiento depende de la coalescencia de las transacciones de memoria.
• Ineficiencia en la construcción: Los métodos de construcción de árboles "top-down" (de arriba hacia abajo) son costosos y difíciles de paralelizar eficientemente en GPU.

El objetivo es desarrollar un marco de trabajo que permita ejecutar búsquedas de vecinos más cercanos (KNN) y agrupamiento "Friends-of-Friends" (FoF) en GPUs con un rendimiento superior a las bibliotecas existentes, aprovechando al máximo el alto rendimiento de cómputo y la eficiencia energética de las GPUs.

2. Metodología

Los autores proponen JZ-TREE, un marco de trabajo implementado en JAX y CUDA que utiliza una jerarquía de árboles basada en orden Morton (orden-z) diseñada específicamente para GPUs.

A. Construcción del Árbol (Bottom-Up)

En lugar de construir un árbol binario anidado profundamente, el método utiliza un enfoque de abajo hacia arriba:

1. Ordenamiento Z-Order: Se ordenan los puntos de entrada utilizando un comparador personalizado que opera sobre coordenadas de punto flotante completo (sin truncar precisión), basado en el bit más significativo que difiere entre dos puntos.
2. Jerarquía de Planos (Tree-Planes): En lugar de nodos binarios, se construyen "planos" de árboles.
   • Se definen nodos como segmentos contiguos en el array ordenado.
   • Se seleccionan puntos de división para crear nodos que contengan un número máximo de puntos ($N_{max}$).
   • Requisito Crítico: Todos los puntos que caen dentro de la misma celda de orden Z deben permanecer juntos en la misma hoja.
   • Tamaño de las Hojas: Las hojas contienen como máximo 48 puntos. Esto significa que el tamaño de las hojas puede variar, pero está acotado superiormente por 48, respetando la integridad de las celdas Z-Order.
   • Se construyen planos más gruesos iterativamente sobre los puntos de división del plano anterior.
   • Resultado: Una jerarquía con profundidad fija y pequeña, donde los hijos de un nodo se almacenan de forma contigua en memoria.

B. Recorrido Dual de Árboles (Dual Tree Walk)

Se implementa un algoritmo de recorrido dual para interactuar entre dos conjuntos de puntos (o un conjunto consigo mismo):

• Listas de Interacción: Se generan listas de interacción que especifican qué nodos de un árbol deben interactuar con nodos de otro.
• Ejecución Colaborativa: Los hilos de la GPU trabajan en grupos para procesar interacciones entre grupos de nodos.
• Coalescencia de Memoria: Gracias a la disposición contigua de los datos en los planos del árbol, los hilos acceden a la memoria global de manera coalescida (patrones regulares), maximizando el ancho de banda.
• Poda Eficiente: Se utilizan distancias inferiores ($d_{low}$) y superiores ($d_{up}$) entre nodos para descartar interacciones innecesarias antes de llegar a los puntos individuales.

C. Implementación en JAX/CUDA

• Se utiliza la interfaz de función extranjera (FFI) de JAX para invocar kernels CUDA.
• Se gestiona la memoria con tamaños estáticos conocidos en tiempo de compilación (JIT) mediante estimaciones heurísticas, permitiendo la compilación just-in-time.
• Soporte para Multi-GPU: Se utiliza un enfoque de particionamiento basado en muestreo para distribuir los datos y mantener la coherencia del orden Z global, minimizando la comunicación entre dispositivos.

3. Contribuciones Clave

1. Nueva Jerarquía de Árbol: Una estructura basada en planos y orden Morton que elimina la divergencia de hilos y garantiza patrones de acceso a memoria altamente coalescidos, superando las limitaciones de los árboles espaciales tradicionales en GPU.
2. Algoritmos Optimizados: Implementaciones eficientes de:
   • K-Nearest Neighbour (KNN): Búsqueda exacta de los $k$ vecinos más cercanos.
   • Friends-of-Friends (FoF): Algoritmo de agrupamiento para identificar componentes conectados (crucial en cosmología).
3. Librería de Código Abierto (JZ-TREE): Una implementación disponible en GitHub y PyPI que integra estas técnicas en el ecosistema JAX, facilitando su uso en aplicaciones científicas modernas.
4. Escalabilidad: Demostración de escalado fuerte en sistemas distribuidos de múltiples GPUs.

4. Resultados

Los autores compararon JZ-TREE con bibliotecas competidoras (SCIPY-CKDTREE, FAISS, CUPY-KNN, JAXKD-CUDA, CLOVER) en hardware NVIDIA A100.

• Rendimiento en KNN:
  • Para tamaños de problema grandes ($N \gtrsim 10^7$), JZ-TREE supera a las bibliotecas competidoras en más de un orden de magnitud (10x a 100x más rápido).
  • Supera a CLOVER (el competidor más cercano) en problemas grandes, donde CLOVER comienza a escalar cuadráticamente.
  • Es significativamente más rápido que las implementaciones basadas en CPU (SCIPY) y otras implementaciones de GPU basadas en KD-trees.
• Rendimiento en FoF:
  • En simulaciones cosmológicas, JZ-TREE es aproximadamente 5 veces más rápido que GADGET4 (32 núcleos CPU), 18 veces más rápido que JFOF (GPU) y hasta 116 veces más rápido que GADGET4 (1 núcleo).
  • Logra calcular catálogos de grupos FoF para $2048^3$ puntos en 64 GPUs en aproximadamente 3 segundos.
• Escalabilidad Multi-GPU:
  • El método muestra un escalado fuerte hasta 64 GPUs. La eficiencia disminuye ligeramente al pasar de un nodo a múltiples nodos debido a la latencia de comunicación, pero el rendimiento general se mantiene alto.
• Robustez: El rendimiento se mantiene consistente a través de diferentes distribuciones de datos (uniforme, normal, distribuciones cosmológicas reales).

5. Significado e Impacto

Este trabajo es fundamental para la comunidad de HPC y astrofísica por varias razones:

• Cambio de Paradigma: Demuestra que los algoritmos basados en árboles, tradicionalmente difíciles de paralelizar en GPU, pueden ser altamente eficientes si se rediseñan desde cero para aprovechar la coalescencia de memoria y la ejecución colaborativa de hilos.
• Aceleración de Simulaciones: Permite realizar evaluaciones repetidas de simulaciones (como inferencia basada en simulaciones) que antes eran prohibitivas por tiempo de cómputo.
• Integración Moderna: Al estar construido sobre JAX, permite combinar el rendimiento de GPU con diferenciación automática y compilación JIT, lo cual es esencial para el aprendizaje automático científico y la optimización de parámetros en simulaciones.
• Aplicabilidad General: Aunque se centra en KNN y FoF, el marco de "planos de árbol" es aplicable a otros algoritmos como métodos multipolo rápidos (FMM) y estimación de funciones de correlación.

En resumen, JZ-TREE resuelve el cuello de botella de la divergencia de hilos y el acceso a memoria en GPU para algoritmos de árbol, ofreciendo una solución práctica y extremadamente rápida para problemas de búsqueda espacial a gran escala.