Track papers across labs, topics, and trends.
Search papers, labs, and topics across Lattice.
This paper introduces JZ-Tree, a novel Morton-order based tree hierarchy optimized for GPU architectures to address thread divergence and irregular memory access patterns common in spatial tree traversal algorithms. By flattening the data layout and enabling collaborative execution across thread groups, JZ-Tree achieves highly coalesced memory access patterns during dual-tree traversal. The authors demonstrate over an order-of-magnitude performance improvement compared to existing GPU libraries for large-scale k-NN search and friends-of-friends clustering, with strong scaling on multi-GPU systems.
Spatial algorithms can be 10x faster on GPUs if you restructure your tree traversal to minimize thread divergence and maximize memory coalescing, as demonstrated by JZ-Tree.
Algorithms based on spatial tree traversal are widely regarded as among the most efficient and flexible approaches for many problems in CPU-based high-performance computing (HPC). However, directly transferring these algorithms to GPU architectures often yields substantially smaller performance gains than expected in light of the high computational throughput of modern GPUs. The branching nature of tree algorithms leads to thread divergence and irregular memory access patterns -- both of which may severely limit GPU performance. To address these challenges, we propose a Morton (z-order)'plane-based tree hierarchy'that is specifically designed for GPU architectures. The resulting flattened data layout enables efficient dual-tree traversal with collaborative execution across thread groups, leading to highly coalesced memory access patterns. Based on this framework we present implementations of two important spatial algorithms -- exact $k$-nearest neighbour search and friends-of-friends (FoF) clustering. For both cases, we observe more than an order-of-magnitude performance improvement over the closest competing GPU libraries for large problem sizes ($N \gtrsim 10^7$), together with strong scaling to distributed multi-GPU systems. We provide an open-source implementation,'JZ-Tree'(JAX z-order tree), which serves as a foundation for efficient GPU implementations of a broad class of tree-based algorithms.