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

This paper introduces JZ-Tree, an open-source JAX and CUDA framework that employs a GPU-optimized Morton z-order plane-based tree hierarchy to overcome thread divergence and memory access inefficiencies, achieving over an order-of-magnitude performance improvement in $k$-nearest neighbour search and friends-of-friends clustering compared to existing GPU libraries.

The Gist

Imagine you are organizing a massive party with 100 million guests (points in space). Your goal is to quickly find the 16 closest friends for every single person at the party, or to group everyone into "cliques" where everyone is standing within arm's reach of someone else.

Doing this on a standard computer (CPU) is like having one very smart but slow librarian who checks every book one by one. It works, but it takes forever.

Doing this on a modern graphics card (GPU) is like having thousands of tiny, super-fast robots working simultaneously. However, there's a catch: these robots are terrible at following complex, branching instructions. If you ask them to "go left if the guest is wearing a hat, but go right if they are holding a drink," they all get confused, stop to wait for each other, and the line moves slowly. This is called thread divergence.

This paper introduces a new way to organize the party, called JZ-TREE, specifically designed so that the thousands of robots can work in perfect unison without getting confused.

The Problem: The "Twisted" Tree

Traditional methods (like KD-trees) build a decision tree that looks like a deep, twisting maze.

• The CPU way: The librarian walks down the maze, turning left or right based on complex rules.
• The GPU problem: When you send 1,000 robots down this maze, some hit a "turn left" sign, and others hit a "turn right" sign. The robots who want to turn left have to wait for the "turn right" group to finish, and vice versa. They also have to run all over the building to grab different books, causing traffic jams.

The Solution: The "Z-Order" Elevator

The authors propose a new way to organize the data using something called a Morton (Z-order) sort.

The Analogy:
Imagine the party is in a giant 3D room. Instead of building a complex maze, you simply ask everyone to line up in a single, long, winding line that snakes through the room like a Z-shaped path (or a snake eating its own tail).

• People standing next to each other in the line are physically close to each other in the room.
• This creates a flat, predictable list rather than a deep, branching tree.

How JZ-TREE Works

Once everyone is in this Z-shaped line, the authors build a "Tree of Planes" (layers of the party).

1. The Layers: Instead of a deep tree, they create flat layers.

   • Layer 0 (The Leaves): These are groups of up to 48 people standing next to each other in the line. Crucially, the system bundles these groups so that anyone standing in the same small "Z-order cell" stays together in the same group, even if it means the group has fewer than 48 people.
   • Layer 1: Groups of those leaves combined into bigger blocks.
   • Layer 2: Even bigger blocks.
   • The Magic: Every layer has the same depth. It's like a stack of pancakes where every pancake is the same thickness. This is predictable for the robots.

2. The Dual Walk (The Team Huddle):
   When the robots need to find neighbors, they don't wander alone. They use a Dual Tree Walk.

   • Imagine two teams of robots looking at two different layers of the party.
   • Because the data is organized in this flat, Z-order way, the robots can grab a whole chunk of data at once (like grabbing a whole shelf of books instead of one book at a time). This is called coalesced memory access.
   • They work together in groups. If one robot finds a potential friend, the whole group checks it instantly. No one is left waiting.

The Results: Speeding Up the Party

The paper tested this on two big problems:

1. K-Nearest Neighbors (KNN): Finding the closest friends.
2. Friends-of-Friends (FoF): Grouping people into cliques.

The Performance:

• Old GPU methods: Were about 10 to 100 times slower than this new method for huge parties (10 million+ people).
• The New Method (JZ-TREE): It scales beautifully. If you add more GPUs (more robots), the speed increases almost perfectly. They can process 100 billion neighbor checks in just a few seconds.

Why This Matters

This isn't just about finding friends at a party. This technology is crucial for:

• Cosmology: Simulating how galaxies form and cluster in the universe.
• Physics: Simulating how millions of particles interact in fluids or gases.
• AI: Making machine learning models faster by quickly finding similar data points.

The Takeaway

The authors took a problem that was hard for GPUs (complex, branching trees) and turned it into a problem GPUs love (flat, predictable, organized lists). By arranging the data in a "Z-shape" and letting the robots work in synchronized teams, they unlocked a massive speedup, making super-complex simulations possible in a fraction of the time it used to take.

They even made the code open-source (called JZ-TREE), so anyone can use this "super-organized party planner" for their own scientific discoveries.

Technical Summary

1. Problem Statement

High-performance computing (HPC) applications are increasingly migrating from CPUs to GPUs due to superior arithmetic throughput and energy efficiency. However, many state-of-the-art algorithms designed for CPUs, particularly those based on spatial tree traversals (e.g., KD-trees, Oct-trees), perform poorly when ported directly to GPUs.

The primary bottlenecks include:

• Thread Divergence: Tree algorithms often involve irregular branching, causing threads within a GPU warp to execute different paths, forcing serialization and reducing efficiency.
• Irregular Memory Access: Conventional tree layouts result in non-coalesced memory access patterns, where threads in a warp access scattered memory addresses, severely limiting bandwidth utilization.
• Inefficiency of Existing GPU Solutions: While brute-force methods offer regular memory access, they lack scalability for large datasets ($N \gtrsim 10^7$). Existing GPU tree methods (like KD-trees) often suffer from expensive construction and poor traversal performance.

The paper addresses the need for a tree-based spatial search framework specifically designed for GPU architectures that minimizes divergence and maximizes memory coalescence.

2. Methodology: JZ-TREE Framework

The authors propose JZ-TREE, a novel framework based on a Morton (z-order) plane-based tree hierarchy. The core innovation is moving away from deeply nested, irregular binary trees toward a structure with fixed depth and contiguous data layouts.

A. Tree Construction

1. Z-Order Sorting: Input points are sorted using a custom comparison operator based on the Most Significant Bit (MSB) of floating-point coordinates. This avoids the precision loss associated with integer Morton encoding while maintaining the spatial locality of the Z-order curve.
2. Bottom-Up Construction: Unlike top-down KD-tree construction, JZ-TREE builds the tree bottom-up from the sorted array.
3. Tree-Planes: Instead of a binary tree, the structure is organized into "tree-planes."
   • Leaf Nodes: Defined by a maximum point count of 48. Crucially, all points within the same z-order cell are kept together in the same leaf. Consequently, leaf sizes vary but are strictly bounded at 48 points.
   • Hierarchy: Coarser planes are generated iteratively by grouping nodes from the finer plane, increasing the max point count by a factor $c$ (default $c=8$).
   • Fixed Depth: The hierarchy has a fixed, small depth, making traversal highly predictable.
   • Contiguous Layout: Children of a node are stored contiguously in memory, enabling fully coalesced memory reads.

B. Dual Tree Walk

The framework implements a dual-tree walk algorithm where interactions between groups of nodes are processed collaboratively.

• Interaction Lists: The algorithm maintains an interaction list of node pairs that need to be checked.
• Pruning: It uses lower ($d_{low}$) and upper ($d_{up}$) distance bounds between node centers and extents to prune unnecessary traversals.
• Kernel Execution: The traversal is implemented via a sequence of CUDA kernels (invoked via JAX's FFI) that move the interaction list from coarse planes to fine planes.
  • FINDRMAX: Calculates the maximum radius required to find $k$-nearest neighbors for a node.
  • COUNT/INSERT: Filters interactions based on distance bounds and updates the interaction list.
  • LEAF-TO-LEAF: Performs the final point-to-point checks within the leaf nodes.

C. Multi-GPU and JAX Integration

• Distributed Execution: The algorithm supports multi-GPU systems by exchanging only the necessary remote node data (splitting points and centroids) at each tree level. Communication is minimized and predictable.
• JAX Compatibility: The implementation leverages JAX for just-in-time (JIT) compilation and automatic differentiation. Memory allocations are estimated statically to satisfy JAX's requirements, with runtime checks for overflow.

3. Key Contributions

1. GPU-Optimized Tree Hierarchy: A novel "plane-based" tree structure that eliminates thread divergence and ensures coalesced memory access, specifically tailored for GPU architectures.
2. Dual Tree Walk Implementation: A highly efficient dual-tree traversal formulation that processes node interactions collaboratively across thread groups.
3. Open-Source Library (JZ-TREE): A production-ready implementation in JAX/CUDA available on GitHub and PyPI, supporting both Exact k-Nearest Neighbor (kNN) search and Friends-of-Friends (FoF) clustering.
4. Scalability: Demonstrated strong scaling to distributed multi-GPU systems (up to 64 GPUs) for problem sizes exceeding $10^8$ points.

4. Results and Performance

The authors benchmarked JZ-TREE against leading libraries (SCIPY-CKDTREE, FAISS, CUPY-KNN, JAXKD-CUDA, CLOVER, GADGET4, HFOF) on NVIDIA A100 GPUs.

• k-Nearest Neighbor (kNN) Search:

  • Performance Gain: For large problem sizes ($N \gtrsim 10^7$), JZ-TREE achieves more than an order of magnitude (10x+) speedup over the closest competing GPU libraries (e.g., CLOVER, JAXKD-CUDA).
  • Comparison: It is ~100x faster than CPU-based KD-trees (SCIPY) and significantly outperforms brute-force approaches (FAISS) at scale.
  • Scaling: Shows near-linear scaling with problem size and excellent weak scaling across multiple GPUs.

• Friends-of-Friends (FoF) Clustering:

  • Performance Gain: On cosmological simulation data ($N \approx 10^8$), JZ-TREE is ~5x faster than GADGET4 (32 CPU cores), ~18x faster than JFOF (1 GPU), and ~116x faster than GADGET4 (1 CPU core).
  • Efficiency: Can generate FoF catalogs for $2048^3$ points on 64 GPUs in approximately 3 seconds.

• Robustness: The method performs consistently across various distributions (uniform, Gaussian, cosmological simulations) and boundary conditions (periodic/non-periodic).

5. Significance

• Paradigm Shift: JZ-TREE challenges the notion that tree-based algorithms are inherently unsuitable for GPUs. By redesigning the data layout and traversal logic, it bridges the gap between the flexibility of tree methods and the throughput of GPUs.
• HPC Impact: The library provides a critical tool for modern scientific simulations requiring massive spatial searches, such as Smoothed Particle Hydrodynamics (SPH), Self-Interacting Dark Matter simulations, and Cosmological Halo Finding.
• Extensibility: The framework is not limited to kNN or FoF; the authors note it serves as a foundation for other tree-based algorithms like Fast Multipole Methods (FMM) and density-based clustering (DBSCAN).
• Software Ecosystem: By integrating with JAX, JZ-TREE enables differentiable physics simulations and seamless integration into modern machine learning and scientific computing workflows.

In conclusion, JZ-TREE represents a significant advancement in GPU computing for spatial problems, offering a highly efficient, scalable, and flexible solution that outperforms existing state-of-the-art methods by an order of magnitude for large-scale datasets.