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

O artigo apresenta o JZ-Tree, uma implementação em JAX e CUDA que utiliza uma hierarquia de árvore baseada em ordem Morton para superar os desafios de divergência de threads e acesso irregular à memória, permitindo buscas de vizinhos exatos e agrupamento "friends-of-friends" com desempenho mais de dez vezes superior ao de bibliotecas concorrentes em GPUs.

O essencial

Imagine que você tem uma biblioteca gigante com bilhões de livros espalhados pelo chão. Se você quiser encontrar os 10 livros mais próximos de um ponto específico, ou agrupar todos os livros que estão "de mãos dadas" (perto uns dos outros), fazer isso um por um seria extremamente lento.

É aqui que entra o JZ-TREE, uma nova ferramenta criada por pesquisadores da Universidade de Viena para computadores superpotentes (chamados de GPUs, os mesmos usados em jogos e inteligência artificial).

Aqui está a explicação do que eles fizeram, usando analogias do dia a dia:

1. O Problema: O Caos na Biblioteca

Os computadores comuns (CPUs) são como bibliotecários muito inteligentes que leem um livro de cada vez, mas fazem isso muito rápido e com muita lógica. Eles usam árvores de dados (como um mapa de organização) para encontrar coisas rápido.

Os computadores gráficos (GPUs), por outro lado, são como milhares de estagiários trabalhando ao mesmo tempo. Eles são incrivelmente rápidos em tarefas repetitivas, mas se você pedir para eles seguirem um mapa de árvore complexo e cheio de desvios (como "se o livro for vermelho, vá para a esquerda; se for azul, vá para a direita"), eles ficam confusos. Alguns estagiários ficam parados esperando, enquanto outros correm. Isso é chamado de "divergência de threads" e desperdiça a velocidade do computador.

2. A Solução: O Mapa do "Z" (JZ-TREE)

Os autores criaram uma nova maneira de organizar esses livros (ou dados) que é perfeita para esses milhares de estagiários.

• A Analogia do Zigue-Zague (Ordem Z): Em vez de fazer uma árvore complexa e profunda, eles organizaram todos os pontos em uma linha reta, seguindo um padrão de "Z" (como o desenho de uma serpentina). Imagine que você tem um mapa de uma cidade e, em vez de listar os endereços por bairro, você lista todos os endereços seguindo um caminho contínuo que cobre a cidade inteira sem pular muito.
• O "Plano" em vez da Árvore: Eles não construíram uma árvore com galhos infinitos. Eles criaram "planos" (camadas). É como se você tivesse uma pilha de panfletos. No topo, você tem grandes grupos de bairros. Abaixo, você tem ruas. Abaixo, casas. Mas, ao contrário de uma árvore onde cada galho pode ter um tamanho diferente, aqui todos os "planos" têm a mesma profundidade. Isso significa que todos os estagiários sabem exatamente o que fazer ao mesmo tempo, sem ficar confusos.
• O Tamanho dos Grupos: Uma regra importante é que cada "caixa" final (folha) contém no máximo 48 pontos. A regra não é que elas tenham exatamente 48, mas sim que, se os pontos estiverem juntos no caminho do "Z", eles devem ficar juntos no mesmo grupo. Isso garante que pontos que são vizinhos no espaço nunca sejam separados, criando grupos de até 48 itens que são sempre espacialmente próximos.

3. Como Funciona na Prática?

O JZ-TREE faz duas coisas principais muito rápido:

1. Encontrar Vizinhos (KNN): Imagine que você quer saber quem são os 10 vizinhos mais próximos de uma casa.

   • O jeito antigo: Você teria que medir a distância de sua casa para quase todas as outras casas.
   • O jeito JZ-TREE: Como os pontos estão organizados no "Z" e agrupados em blocos de até 48, o computador pode olhar para blocos inteiros de casas de uma vez. Se um bloco inteiro está muito longe, ele ignora o bloco todo. Se está perto, ele olha dentro do bloco. Como todos os estagiários olham blocos organizados lado a lado, a memória do computador é lida de forma super eficiente (como ler uma fila de pessoas em vez de pular de uma para outra).

2. Agrupar Amigos (Friends-of-Friends - FoF): Imagine que você quer formar grupos de amigos onde todos estão a menos de 1 metro de distância um do outro.

   • O JZ-TREE usa a mesma lógica. Ele conecta os pontos que estão perto, formando "ilhas" de amigos. Isso é crucial para simulações de cosmologia (estudar como galáxias se formam), onde bilhões de partículas precisam ser agrupadas.

4. O Resultado: Velocidade Insana

O papel mostra que essa nova organização é mais de 10 vezes mais rápida do que as melhores ferramentas atuais para problemas grandes (com milhões ou bilhões de pontos).

• Escala: Eles testaram em um único computador e em 64 computadores trabalhando juntos. Funciona perfeitamente em ambos.
• O "Pulo do Gato": A grande vantagem é que eles conseguiram fazer algoritmos que normalmente são "chatos" para GPUs (porque exigem muita lógica de "se/então") rodarem de forma super organizada, como se fosse uma linha de montagem.

Resumo em uma frase

Os pesquisadores criaram um novo "mapa de organização" (chamado JZ-TREE) que transforma problemas complexos de encontrar vizinhos e agrupar dados em tarefas simples e lineares, permitindo que supercomputadores gráficos resolvam problemas que antes levavam horas, em apenas segundos.

É como transformar um labirinto confuso em um corredor reto e largo, onde milhares de pessoas podem correr juntas sem bater umas nas outras.

Resumo técnico

Resumo Técnico: JZ-TREE

1. O Problema

Algoritmos baseados em árvores espaciais (como árvores KD) são padrão em computação de alto desempenho (HPC) baseada em CPU para problemas de busca de vizinhos e agrupamento. No entanto, a transferência direta desses algoritmos para arquiteturas de GPU frequentemente resulta em ganhos de desempenho muito menores do que o esperado. As principais causas são:

• Divergência de Threads: A natureza ramificada da travessia de árvores faz com que threads dentro de um warp (grupo de execução na GPU) sigam caminhos diferentes, forçando a serialização da execução.
• Padrões de Acesso à Memória Irregulares: A estrutura de árvore tradicional leva a acessos à memória global não coalescidos, o que é um gargalo crítico no throughput das GPUs.
• Ineficiência em Grandes Escalas: Métodos existentes muitas vezes não escalam bem para grandes volumes de dados ($N \gtrsim 10^7$) em sistemas distribuídos multi-GPU.

O objetivo deste trabalho é criar uma estrutura de dados e algoritmos de travessia otimizados especificamente para GPUs, mantendo a precisão exata (sem aproximações) e aproveitando o alto paralelismo das GPUs modernas.

2. Metodologia

Os autores propõem o JZ-TREE, um framework baseado em uma hierarquia de árvores de plano (tree-planes) construída sobre uma ordenação Morton (ordem-z).

• Construção da Árvore (Bottom-Up):

  • Em vez de construir uma árvore binária profunda de cima para baixo (top-down), o método primeiro ordena os pontos usando a ordem-z (Morton).
  • A árvore é construída de baixo para cima, criando uma hierarquia de "planos" onde cada nível contém nós que particionam o espaço.
  • Estrutura de Planos: Diferente de árvores binárias tradicionais, esta hierarquia possui profundidade fixa e permite que os nós tenham um número variável de filhos. Os filhos de um nó são armazenados contiguamente na memória, permitindo leituras totalmente coalescidas.
  • Regularização e Tamanho das Folhas: Para evitar nós excessivamente grandes em regiões de baixa densidade (comuns em distribuições gaussianas ou simulações cosmológicas), é aplicado um critério de regularização. As folhas contêm no máximo 48 pontos, com a restrição fundamental de que todos os pontos dentro da mesma célula de ordem-z devem permanecer juntos na mesma folha. Isso significa que os tamanhos das folhas variam, mas são estritamente limitados a 48, preservando a coesão espacial local.
  • Multi-Tipos de Pontos: O framework suporta múltiplos tipos de pontos (ex: pontos de consulta e fonte distintos) em uma única árvore, concatenando-os antes da ordenação e rastreando contagens separadas durante a construção.

• Algoritmo de Travessia Dual (Dual Tree Walk):

  • Utiliza uma abordagem de travessia dual para processar interações entre grupos de nós colaborativamente.
  • Listas de Interação: Em vez de verificar pares ponto-a-ponto imediatamente, o algoritmo gera listas de interação entre nós (leaf-to-leaf) que são refinadas de cima para baixo (do nível raiz até as folhas).
  • Otimização de Memória: A estrutura permite que threads dentro de um bloco carreguem dados de nós irmãos na memória compartilhada (shared memory) de forma colaborativa, minimizando acessos à memória global.
  • Poda Eficiente: Utiliza limites de distância inferior ($d_{low}$) e superior ($d_{up}$) para podar interações desnecessárias antes de chegar aos pontos individuais.

• Implementação em JAX/CUDA:

  • A implementação é feita em JAX, utilizando kernels CUDA via interface de função estrangeira (FFI).
  • Projetado para compilação just-in-time (JIT), com alocações de memória estimadas em tempo de compilação para garantir estabilidade e performance.

3. Contribuições Principais

1. Novo Framework de Árvore para GPU: Uma hierarquia baseada em planos e ordenação Morton que elimina a divergência severa de threads e maximiza a coalescência de memória, respeitando a coesão de células Morton.
2. Implementação de Algoritmos Críticos:
   • Busca Exata de k-Vizinhos Mais Próximos (kNN): Implementação exata (não aproximada) que supera bibliotecas concorrentes.
   • Agrupamento Friends-of-Friends (FoF): Implementação eficiente para identificar componentes conectados em grafos espaciais, crucial para cosmologia.
3. Escalabilidade Multi-GPU: O algoritmo foi adaptado para sistemas distribuídos, onde cada GPU mantém nós locais e solicita dados remotos apenas quando necessário, mantendo a comunicação previsível e baixa.
4. Código de Código Aberto: Disponibilização da biblioteca JZ-TREE no GitHub e PyPI, servindo como base para outros algoritmos baseados em árvores (ex: Métodos Multipolo Rápidos, DBSCAN).

4. Resultados e Desempenho

Os testes foram realizados no cluster Leonardo (CINECA) com GPUs NVIDIA A100.

• Comparação com kNN:

  • Para problemas grandes ($N \gtrsim 10^7$), o JZ-TREE supera as bibliotecas concorrentes (FAISS, CLOVER, KD-tree em CPU e GPU) em mais de uma ordem de magnitude (10x a 100x mais rápido).
  • Em comparação com o CLOVER (o concorrente mais próximo em GPU), o JZ-TREE é significativamente mais rápido em grandes escalas, onde o CLOVER começa a escalar quadraticamente.
  • A construção da árvore e a ordenação consomem apenas cerca de 20% do tempo total, com a maior parte do tempo sendo gasta nas interações folha-para-folha.

• Escalabilidade Multi-GPU:

  • O método escala bem de 1 a 64 GPUs.
  • A eficiência cai ligeiramente ao passar de 1 para 2 GPUs (devido à comunicação inicial), mas a degradação adicional ao escalar até 64 GPUs é de apenas ~30%.
  • Foi capaz de calcular vizinhos para $10^{11}$ pares de pontos em cerca de 1,3 segundos usando 64 GPUs.

• Agrupamento FoF:

  • Em simulações cosmológicas, o JZ-TREE-FOF foi 5x mais rápido que o GADGET4 (32 núcleos CPU), 18x mais rápido que o JFOF (GPU) e 116x mais rápido que o GADGET4 (1 núcleo).
  • Conseguiu gerar catálogos de grupos para $2048^3$ pontos em 64 GPUs em cerca de 3 segundos.

• Robustez: O desempenho é consistente em diferentes distribuições de dados (uniforme, Gaussiana, simulações cosmológicas reais) e condições de contorno periódicas.

5. Significado e Impacto

Este trabalho representa um avanço significativo na aplicação de algoritmos baseados em árvores em GPUs. Ao resolver os problemas fundamentais de divergência de threads e acesso à memória, o JZ-TREE permite que problemas que antes eram limitados a CPUs ou requeriam aproximações grosseiras sejam resolvidos com precisão exata e alta velocidade em GPUs.

• Aplicações: É particularmente relevante para simulações de hidrodinâmica de partículas (SPH), simulações de matéria escura auto-interagente e busca de halos em cosmologia.
• Futuro: O framework estabelece uma base sólida para implementar outros métodos complexos em GPU, como o Método Multipolo Rápido (FMM), prometendo acelerar drasticamente a próxima geração de simulações científicas de alto desempenho.

Em resumo, o JZ-TREE demonstra que, com o design correto de estrutura de dados e padrões de acesso à memória, é possível extrair o máximo potencial das GPUs para problemas de busca espacial exata, superando em muito as soluções atuais de CPU e GPU.