GPU-FRIENDLY KNN AND FOF WITH DUAL TREE WALKS IN JAX/CUDA
1
JZ - TREE: GPU friendly neighbour search and
friends-of-friends with dual tree walks in JAX plus CUDA
arXiv:2604.05885v1 [cs.DC] 7 Apr 2026
Jens Stücker
, Oliver Hahn
, Lukas Winkler
Abstract—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 ≳ 107 ), together with strong scaling to distributed multiGPU 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. Index Terms—GPU computing, nearest neighbour search, friends-of-friends, high performance computing
I. I NTRODUCTION IGH-performance computing (HPC) applications are increasingly shifting from CPU-based implementations to graphics processing units (GPUs). This shift is motivated both by the high arithmetic throughput and by the favorable energy efficiency of GPUs, which typically provide substantially more floating-point operations per unit power than conventional CPUs. Further, the reduction in execution time enables classes of applications that require not just a single large simulation, but a large number of repeated evaluations – for example simulation-based inference [1]. In addition, recent software frameworks such as JAX make it possible to combine accelerator-based performance with justin-time compilation, automatic differentiation, and a highlevel programming model, which is particularly attractive for modern scientific applications [2]. However, GPUs differ fundamentally from CPUs in how performance is achieved. GPUs follow a throughput-oriented
H
All authors are at the University of Vienna, Department of Astrophysics, Türkenschanzstraße 18, 1180 Vienna, Austria. Oliver Hahn and Thomas Flöss are additionally at University of Vienna, Department of Mathematics, OskarMorgenstern-Platz 1, 1090 Vienna, Austria. E-mail: [email protected]
, Adrian Gutierrez Adame
, and Thomas Flöss
Uncoalesced Access Thread T0
Thread T1
Thread T2
Coalesced Access
Thread T3
Thread T0
Thread T1
Thread T2
Thread T3
6
9 10 11
step 0 step 1 0
1
2
3
4
5
6
7
8
9 10 11
0
1
2
3
4
5
7
8
Global Memory Fig. 1. Illustration of uncoalesced versus coalesced memory access. Coalesced memory access patterns achieve significantly better performance on GPU than uncoalesced access.
parallel execution model with large numbers of lightweight threads, which differs significantly from the latency-optimized design of CPUs [3]. As a consequence, many algorithms that are state-of-the-art on CPUs perform poorly when transferred directly to GPUs without redesign. Efficient GPU implementations typically require minimizing host-device communication, reducing global memory traffic, limiting thread divergence, and maximizing memory coalescence. In particular, coalesced memory access (illustrated in Figure 1) is a primary performance consideration, as memory transactions are shared across threads within a warp [4]. Similarly, divergent control flow within a warp can significantly degrade performance, since threads executing different branches must be serialized [3]. In practice, these constraints favor algorithms with regular control flow, predictable memory access patterns, and a limited number of synchronization points. A. Related Work Tree-based data structures are a particularly important example of this challenge. On CPUs, trees are a standard tool for reducing the complexity of spatial search and interaction problems [5], [6]. They are used in nearest-neighbour search [7], friends-of-friends clustering, N -body methods [8], multipole schemes [9], and many related algorithms. Yet on GPUs, tree methods are often much less competitive than their asymptotic complexity would suggest. Tree construction is frequently expensive, traversal tends to induce thread divergence, and the associated memory access patterns are often highly irregular [10]–[12]. While iterative traversal schemes can reduce some forms of control-flow divergence [13], they generally do not eliminate divergence in the number of traversal steps taken by different threads. Further, conventional tree layouts make it difficult for neighbouring threads to read memory
GPU-FRIENDLY KNN AND FOF WITH DUAL TREE WALKS IN JAX/CUDA
collaboratively, so that even moderate divergence in traversal may quickly destroy memory coalescence. In particular, nearest neighbour search has been studied extensively, and a wide range of algorithmic approaches have been proposed. Classical exact methods are typically based on spatial tree structures such as KD-trees [7] or ball trees, which enable logarithmic query complexity in low dimensions. Variants of dual-tree traversal further improve efficiency by processing interactions between groups of nodes jointly [14], [15]. Closely related approaches based on uniform grids or cell lists are widely used in particle simulations, where the domain is decomposed into regular bins to enable efficient neighbour queries with predictable memory access patterns [16], [17]. On modern hardware, particularly GPUs, brute-force approaches based on dense distance evaluations have become increasingly competitive due to their regular memory access patterns and high arithmetic intensity [18]–[20]. In addition, a large body of work has focused on approximate nearest neighbour (ANN) methods, including hashing-based techniques and product quantization [21], as well as graph-based approaches such as nearest-neighbour graphs and navigable small-world structures [22]. These methods often achieve significantly improved query times at the cost of approximation error. Finally, a notable recent advancement for exact neighbour search on GPUs is CLOVER [23], a spatio-graph-based method that constructs an index of random Voronoi partitions to prune the search space while maintaining high hardware utilization, outperforming prior tree methods by an order of magnitude for some setups. B. Contributions In this work we propose a novel tree framework designed specifically to address the constraints of GPU architectures. The presented hierarchy is based on Morton, or z-order, sorting and can be constructed efficiently in a bottom-up fashion. Rather than producing a deeply nested binary tree with irregular traversal depth, the construction yields a hierarchy of tree-planes with fixed and small depth. This makes tree walks highly predictable and allows them to be implemented through a small number of kernel launches. In addition, the hierarchy is organized such that the children of a node are stored contiguously and may be accessed with fully coalesced memory reads. Combined with a dual tree walk formulation, this allows interactions between groups of nodes to be processed collaboratively, reducing redundant memory access compared to more conventional traversal schemes. We demonstrate the benefits of the framework with two algorithms, k-nearest neighbour search and friends-of-friends (FoF) clustering. For both cases, we find significant performance improvements over the closest competiting libraries, reaching more than an order of magnitude improvement for sufficiently large problems. The presented framework is not specific to these two use cases. The same tree representation and traversal strategy can be extended naturally to a range of other tree-based algorithms, including density-based clustering methods such as DBSCAN, fast multipole methods, and correlation function estimation.
2
Our main optimization target is for low dimensions (d ∼ 3), and large point counts N ≫ 106 – a regime that is highly relevant for many HPC simulation codes. We are less concerned with very high-dimensional settings or with small problem sizes, although we will show that the presented methods remain competitive outside of the primary target regime as well. Here, we only consider a Euclidean distance measure, but including other distance measures in the future would be viable. The remainder of this paper is structured as follows. In Section II we introduce the z-order based tree construction. We then describe the nearest neighbour algorithm and evaluate its performance in Section III and the FoF algorithm in Section IV. Finally, we conclude in Section V. The publication is accompanied by an open-source implementation of the presented algorithms named ’JZ - TREE’ (short for JAX z-order tree) that is available on GitHub1 and PyPI2 with additional documentation and usage examples under3 . II. Z- ORDER TREES We construct a tree in two steps: (1) A sort of the input position array in Morton / z-order [24] and (2) A search for splitting points on the position array to summarize points (and nodes) into coarser nodes. Both steps involve only GPU friendly operations on flat arrays so that the tree construction is significantly faster on GPUs than widely used top-down construction methods of KDTrees. We note that Peano-Hilbert (PH) order is often considered superior to z-order in terms of spatial locality [25], [26]. However, we prefer z-order here due to its simplicity and flexibility. In particular, z-order can be defined directly for all floating-point coordinates without requiring a predefined domain or refinement level. In contrast, PH order is typically constructed on discretized grids and becomes more complex to generalize across dimensions or to arbitrary floating-point data. A. Z-order sort The most common and fastest approach to sorting position vectors in z-order is to sort by an integer key obtained by interleaving the bits of the coordinate components (Morton encoding [24], [27]). However, for floating-point positions, defining such a key requires restricting the domain and truncating precision. An alternative approach is to define a custom comparison operator that directly compares full position vectors and to use a sorting algorithm that supports custom comparators, such as mergesort [28]. Here, we adopt this approach, as it provides maximum generality – allowing the construction of tree structures at full floating-point precision. As we show later, the associated performance overhead is negligible, since sorting is not a bottleneck in the presented algorithms. To define this comparison operator, let us assume that we have a function msbfixed available that extracts the most 1 https://github.com/jstuecker/jztree/ 2 https://pypi.org/project/jztree 3 https://jstuecker.github.io/jztree/
GPU-FRIENDLY KNN AND FOF WITH DUAL TREE WALKS IN JAX/CUDA
3
1.0
y
0.5 0.0 0.5 1.0
1.0
0.5 0.0 x
0.5
1.0 1.0
0.5 0.0 x
0.5
1.0
Fig. 2. The z-order curve for a regular grid (left) and for randomly distributed points (right). Colour encodes the linear index in the sorted array.
significant differing bit of two positive fixed-point numbers (normalized to exponent 0). For example, for two numbers a and b common z }| { a = ...1010 1001.0101.. b = ...1010 0101.1011... ↑ most significant differing bit at 23 the most significant differing bit is the first bit at which the two representations differ after their common prefix. We label bits by the power of two that they represent so that msbfixed would return 3 in the example above. Such a function can easily be implemented by counting leading zeros on a bitwise exclusive or of a and b. Given the function msbfixed we may define a more general function msb that acts on floating point numbers to extract the most significant bit that would differ if they were normalized as fixed point numbers with exponent 0. Given a separation into sign, exponent and mantissa: a = s(a) · 2e(a) · m(a)
(1)
In Figure 2 we show two examples of the z-sorted points on a regular grid (left) and for a uniform random distribution (right). For a regular grid, the traversal follows the characteristic Z-shaped (Morton) pattern [24]. In JZ - TREE the z-order sort is implemented via a library call to the mergesort routine of the CUB library [30]. For multiGPU execution, we employ a sampling-based partitioning approach [31]. In a first step, a subset of Nsamp ∼ 1000 points is sampled on each GPU, collected, sorted on a single device and broadcasted. Based on the sorted samples, a set of splitters is chosen such that the sampled points are evenly partitioned. Subsequently, all points are redistributed across GPUs according to these splitters and sorted locally, resulting in a globally consistent z-order. Assuming no duplicate keys, the expected relative imbalp ance scales as O( log(NGPU )/Nsamp ) [32]–[34], leading to imbalances well below 10% in all scenarios that we consider here. After sorting, domain boundaries can be adjusted for perfect load balance. B. Nodes We define a node with center c and Morton level lvl as the set of all points whose hypothetically interleaved binary representations share the leading bits up to lvl with c. Such a node corresponds to a contiguous segment in z-order. Given two points p, q ∈ Rd , let k denote the dimension in which they differ most significantly in the Morton sense. The corresponding Morton level is lvl(p, q) = (msb(pk , qk ) + 1) · d − k.
(5)
where the offset by one accounts for the fact that the common interval is larger than the position of the highest differing bit. We may further define per-dimension extent levels ( 1 if i ≥ (lvl mod d), lvl + (6) li = d 0 otherwise.
where s ∈ {−1, 1}, e ∈ N and m ∈ [1, 2), we may write if s(a) ̸= s(b) EMAX so that Li = 2li corresponds to the spatial extent of the msb(a, b) = max(e(a), e(b)) else if e(a) ̸= e(b) node in the ith component and 2lvl corresponds to the volume of a node. Extent levels may differ at most by one across e(a) + msbfixed (m(a), m(b)) else (2) dimensions, so that nodes can be rectangular with axis ratios of at most two. where EMAX is one larger than the maximum exponent value (e.g., 128 for float32 and 1024 for float64 in IEEE 754 standard [29]). The function msb can be implemented efficiently C. Plane based tree-hierarchy using bitwise operations (bit shifts and leading-zero counting), As a first step to constructing a tree hierarchy we calculate with special care required to handle subnormal numbers, where lvli = lvl(xi−1 , xi ) (7) the mantissa representation differs from normalized values. We can then define the z-order comparison operator of two for each consecutive pair of points i − 1 and i in the sorted vectors p, q ∈ Rd as a comparison along the dimension with array. To simplify later calculations, we assume the existence the largest most significant bit difference: of an additional vector with −∞ components at index −1 k = argmaxi msb(pi , qi ) (3) and +∞ components at index N , so that we obtain N + 1 values for N points. It is useful to interpret these level values p ≺z q ⇔ pk ≤ qk (4) as being associated with the gaps between consecutive points where argmax selects the first occurrence of the maximum – (see Figure 3). so that differences in earlier coordinates are more significant As a next step, we determine the range of points contained than those in later coordinates. in the smallest node that includes points i − 1 and i. To this
GPU-FRIENDLY KNN AND FOF WITH DUAL TREE WALKS IN JAX/CUDA
x
(−∞) 1.6
lvl
(129)
2
−1
3
1
2
4
0
(129)
n
(8)
3
2
6
2
3
8
2
(8)
0
4.6
1 Leaf 0
spl(1) (n > 4)
3.3
5.6
3 Leaf 1
0
5 Leaf 2
2 Node 0
6.8
9.4
6 Leaf 3
(∞)
Leaf 4
Nmax = 4
1.0
Nmax = 8
0.5
8
4 Node 1
9.7
y
spl(0) (n > 2)
3.1
4
0.0
5
0.5
Node 2
1.0
Fig. 3. Illustration of the steps involved in building a split based tree structure. By choosing splitting points that have n > Nmax we obtain tree-planes where all nodes have n ≤ Nmax points.
7
end, we perform a binary search to the left to find the smallest index lb that would be part of such a node lvl(xlb , xi ) ≤ lvli , and a binary search to the right to find the smallest index rb that would be outside lvl(xi−1 , xrb ) > lvli . If no such indices exist, we set lb = 0 and rb = N . The number of points contained in the node is then n = rb − lb . The information contained within lb and rb is in principle sufficient to define a full binary tree, where the parent of each node is given by lb or rb depending on which one has the lower level argmin(lvllb , lvlrb ). However, walking such a binary tree on GPU architectures would lead to poor memory coalescence, since different threads may access very different locations in memory. We therefore choose a different approach here, where we allow nodes to have a variable number of children, but keep the depth of the resulting tree fixed. We define a tree-plane as a set of nodes that partitions the points x, such that each point belongs to exactly one node (while empty regions of space may remain uncovered). A tree-plane may be parameterized through a set of Nnodes+1 splitting points spl in the z-order index space so that a node i contains all points in the range [spli , spli+1 ). Recall that ni is the number of points that would need to be included in a node that contains points i−1 and i. We construct a tree-plane by selecting all separation points with n > Nmax . Intuitively, each tree-plane partitions the points into the largest possible Morton cells subject to the constraint that each cell contains at most Nmax points. In Figure 3 we illustrate the splitting (0) points spl(0) of leaf nodes created this way with Nmax = 2 which we refer to as the ’0th plane’. We may construct coarser tree-planes iteratively by applying the same procedure to the splitting points of the previous plane. That is, we retain only those splits between nodes of (p+1) plane p for which n exceeds Nmax . In Figure 4 we show an (0) example of two tree-planes that are obtained with Nmax = 4 (1) and Nmax = 8 for a uniform random distribution on a twodimensional domain in the range [−1, 1] with N = 100 points. Note a few properties that are different between this tree-plane hierarchy and conventional space-filling binary trees: • Space that doesn’t contain points may or may not be part of a tree-node.
1.0
0.5 0.0 x
0.5
1.0 1.0
6
5
4 lvl
0.5 0.0 x 3
2
0.5
1.0 1
Fig. 4. Nodes that are obtained by selecting the largest Morton nodes that hold at most Nmax = 4 (left) or Nmax = 8 (right) points. These nodes only fill the space that contains points and some nodes may only contain a single points leading to zero extent (marked with red squares).
Some nodes may only contain a single point and have zero extent. • Nodes on the coarser plane may contain a flexible number of nodes of the finer plane. • Some nodes on the coarser plane may have themselves as their own only child on the finer plane. • Different children may have different extent. • The tree has the same fixed depth everywhere. In practice we build a tree-plane hierarchy by choosing (0) Nmax to define leaf nodes for the finest level of the tree and then successively increasing by factor c per plane level: •
(p) (0) p Nmax = Nmax c
(8)
(0) with defaults Nmax
= 48 and c = 8 which we find to be good choices performance wise. If we wanted to end up with (p) a single root node, we could keep coarsening until Nmax ≥ N at which point we’d be guaranteed to have a single node that covers all points. However, on GPUs it is preferable to have a coarsest level that has already a notable number of nodes so that most streaming multiprocessors have work to do from the beginning. We define a target number Ntarget (typically of order 1000) that we aim to obtain at the coarsest level. We may get a rough estimate of the number of nodes that might be contained on a tree-plane with Nmax based on the heuristic that typical nodes should contain at least Nmax /2 points, since otherwise they might be merged with one of their neighbours: Nnodes ≲
N Nmax /2
(9)
This typically overestimates the number of nodes, but it is not a strict upper bound, since in z-order a single high-n node may block multiple low-n nodes from merging. We stop coarsening when the estimated number of nodes is smaller than Ntarget . In JZ - TREE the distributed tree-construction is implemented in four steps: (1) We first adjust the domain to ensure that no node of the coarsest tree-plane crosses domain boundaries. This is achieved by determining how far a node at a given Morton level would extend across the domain boundary. The
GPU-FRIENDLY KNN AND FOF WITH DUAL TREE WALKS IN JAX/CUDA
5
domain then needs to be adjusted to the starting or end point of (last) the largest node that contains ≤ Nmax points. The subsequent tree construction can be treated fully locally from this point (0) on. (2) We extract leaf splitting points Nmax by checking (0) where splitting points exceed Nmax through a range search (0) of Nmax to the left and right of each splitting point. This step is optional, but tends to be faster for the leaf level than a binary search. (3) We then determine n for each leaf splitting point through the earlier described binary search and (4) Extract splitting points for the full hierarchy.
construction, we track the point counts of each type separately for every candidate node. Splitting points are then chosen such (p) that the maximum count over all types does not exceed Nmax for any node on tree-plane p. After construction, points are separated again into typespecific arrays while preserving z-order. Only the leaf-level splits spl(0) are defined separately for each type. This enables coalesced memory access within each species while maintaining a shared tree structure that adapts to all point distributions and keeps the tree traversal simple.
D. Regularization
F. JAX implementation details
For many problem setups (e.g. uniform random distributions or particle distributions from cosmological simulations), the described tree structure is sufficient. However, for distributions that contain a small number of points far from the bulk – e.g. multivariate Gaussian distributions – summarizing nodes solely based on the number of contained points may produce a small number of nodes with very large extent. This is problematic for nearest-neighbour search, where at least a region of size comparable to the node must be explored, which in the worst case can include almost all points. To improve performance in such scenarios, we introduce a simple regularization criterion. For each tree-plane p, we define a global maximum level lvl(p) max and retain all splitting points whose level satisfies lvl > lvl(p) max . Intuitively, this prevents the formation of excessively large nodes in low-density regions by enforcing a global upper bound on node size. We (p) (p) define this maximum level so that the volume Vmax = 2lvlmax of nodes never exceeds
Most computationally intensive parts of our implementation are realized as CUDA kernels, invoked via the foreign function interface (FFI) of JAX. To maintain compatibility with JAX’s just-in-time (JIT) compilation, all memory allocations must have statically known sizes at jit-compile time. Since the number of nodes per tree-plane is data-dependent, we allocate one contiguous buffer for each node property (e.g. splitting points, particle counts, node centers, node levels) and store all tree-planes within this buffer using data-dependent offsets. The required allocation size is estimated as
(p)
(p)
(p) Vmax = 2lvlmax = fmax V90%
allocation size = alloc fac nodes ·
N (0) Nmax /2
,
(11)
where typically alloc fac nodes ∼ 1–2 is sufficient in practice. If the allocated size is insufficient, a runtime error is raised indicating the required increase.
(10) (p)
where we typically choose fmax ∼ 50. Here, V90% denotes the point number weighted average volume of nodes on plane p, computed over the subset of smallest nodes that together contain 90% of all points. This excludes a small number of very large nodes that may otherwise dominate the average. For the scenarios considered in this work, this simple regularization scheme is sufficient. However, we leave the possibility of incorporating more sophisticated techniques in JZ - TREE for future work.
III. N EAREST NEIGHBOUR SEARCH We describe how to implement a k-nearest neighbour search based on the plane-based tree hierarchy that we have described in Section II. The neighbour search happens conceptually in two steps: (1) A dual tree walk on the tree hierarchy to determine per leaf an interaction list of other leaves that need to be checked to guarantee that all candidate neighbours required for an exact k-nearest neighbour search are considered. (2) A neighbour search that traverses the leaf-leaf interaction list collaboratively among points in the same leaf.
E. Multiple point types Some algorithms require treating multiple point types separately in the tree. For example, in nearest-neighbour search, one may wish to query the tree using a set of query points xquery distinct from the source points x. The most common approach is to construct a separate tree for the query points and to treat query and source trees explicitly during the dual tree walk [14], [15], [35]. However, this increases implementation complexity and may lead to suboptimal refinement, as the query tree is constructed independently of the source distribution. Instead, we construct a single tree jointly over all point types. This is achieved by concatenating the positions of all types into a single array prior to the z-order sort. During tree
A. Interaction lists We parameterize an interaction list as a tuple ilist = (isrc, ispl) of two arrays: a set of source indices isrc and a set of splitting points ispl. The interaction list is sorted by receiving nodes so that a receiving node i has to interact with the isrc indices in the range from ispli up to ispli+1 − 1. A dense interaction list where every node out of Nnodes interacts with every other node can be initialized as ispli = Nnodes · i
(12)
isrcj = j mod Nnodes
(13)
2 for i ∈ [0...Nnodes ] and j ∈ [0...Nnodes ].
GPU-FRIENDLY KNN AND FOF WITH DUAL TREE WALKS IN JAX/CUDA
6
Algorithm 1 Dual tree walk for nearest neighbour search
Algorithm 3 Conceptual outline for F IND R MAX kernel. (p)
Require: For each plane p = 0, . . . , P −1: node splits spl , Require: par i, (isrc, ispl) = ilist, spl(p+1) , n(p) , c(p) , source point count n(p) , node centers c(p) , and Morton lvl(p) (p) (0) levels lvl . Source positions x and leaf splits spl . 1: for group step node i = spl[par i] up to spl[par i+1]−1 Query positions xq and splits spl(0) . do q (P ) 2: read node data of node i into registers 1: spl ← R ANGE(0, Ntopnodes , NGR) 3: H ← empty neighbour heap with counts 2: ilist ← D ENSE I NTERACTION L IST(⌈Ntopnodes /NGR⌉) 4: for int = ispl[par i] up to ispl[par i + 1] − 1 do 3: for p = P − 1 down to 0 do par j ← isrc[int] 4: ilist ← N ODE T O N ODE(ilist, spl(p+1) , n(p) , c(p) , lvl(p) ) 5: 6: read node data in range spl[par j]...spl[par j+1]−1 5: end for (0) (0) into shared memory 6: return L EAF T O L EAF(ilist, spl , x, splq , xq ) 7: for node j = spl[par j] up to spl[par j + 1] − 1 do 8: r ← dup (node i, node j) Algorithm 2 Node to Node interaction function. 9: if r < R ADIUS O F C OUNT(H, k) then Require: Interaction list ilist, node splits spl(p+1) , point 10: insert (r, n[node j]) into H count n(p) , centers c(p) , and Morton levels lvl(p) . 11: end if (p+1) 1: Rmax ← F IND R MAX(ilist, spl , n(p) , c(p) , lvl(p) ) 12: end for (p+1) (p) 2: icount ← C OUNT(ilist, spl , c , lvl(p) , Rmax ) 13: end for 3: ispl ← C UMULATIVE S UM P REP 0(icount) 14: Rmax [node i] ← R ADIUS O F C OUNT(H, k) (p+1) (p) 4: isrc ← I NSERT(ilist, spl , c , lvl(p) , Rmax , ispl) 15: end for 5: return Interaction list (isrc, ispl)
B. Dual Tree Walk We sketch the necessary steps for the dual tree walk in Algorithm 1. As a first step, we group the top level nodes into ’pseudo’ super nodes where NGR denotes the grouping size. This grouping is necessary because an entry in the interaction list represents interactions between all children of the receiving node and all children of the source node. Grouping ensures that this assumption remains valid at the top level. A dense interaction list is then initialized on these super nodes so that effectively every top-node will interact with every other topnode. The precise value of NGR is not critical and we typically choose NGR = 32. Subsequently, we evaluate a node-node interaction function on every plane to move the interaction list from plane p + 1 to p and finally evaluate the leaf-leaf interaction list. Given two nodes with centers c1 and c2 and per-dimension extent vectors L1 and L2 (that may be calculated from the Morton level lvl), we can define a lower distance dlow and an upper distance dup as s±,i (a, b) = max(|ai | ± bi , 0) L1 + L2 dlow = s− c1 − c2 , 2 L1 + L2 dup = s+ c1 − c2 , 2
(14) (15) (16)
It is guaranteed that every point in node 1 includes all points from node 2 at a radius Rmax ≥ dup . Further, it is guaranteed that no point of node 2 lies within a radius smaller than dlow from any point in node 1. We can therefore use dup to find guaranteed upper bounds for the radius in which neighbours need to be checked and dlow for efficient pruning of interactions. The node-to-node interaction function is sketched in Algorithm 2. It works in four steps, each of which requires
a separate kernel launch: (1) Determine for each node a maximum radius Rmax that guarantees that it contains the k-th nearest neighbour of all points inside the node. (2) For each node, count the number of nodes for which dlow ≤ Rmax . (3) Calculate the cumulative sum (and prepend 0). (4) Insert the interaction source indices using ispl as relative offsets in the array. Steps (1), (2) and (4) share very similar traversal logic, so we will only discuss step (1) in detail to highlight how the presented data structures can be used to define a CUDA kernel with a good memory access pattern. The prefix sum in step (3) can be implemented through a library call to CUB. The kernel for determining Rmax is outlined in Algorithm 3. Each thread block is assigned a parent node pari , determined by the CUDA block index. The outer loop assigns child nodes nodei of the parent pari to individual threads. If the number of children exceeds the number of threads, multiple iterations are required. Subsequently, the interaction list is traversed over source parent nodes par j. To minimize global memory access, the data of all child nodes of parj is loaded collaboratively into shared memory. Finally, the loop in line 7 iterates over all child nodes (for each thread) to insert the upper node-node distance into the neighbour heap H. The heap data structure H is implemented fully in registers – following [13]. It keeps track of a static number of Nr distances and counts. R ADIUS O F C OUNT gives a preliminary estimate of Rmax as the smallest radius for which the cumulative count exceeds or equals k. If the total count is smaller than k, this estimate is set to ∞. New entries are inserted into the heap to maintain order, discarding the last element. However, if discarding the last element would lead to the heap holding a total count smaller than k, then we instead add the new count to the first element with larger radius. The memory access pattern of the F IND R MAX kernel is ideal for GPU architectures: The global memory accesses in line 2 and line 6 are perfectly coalesced. Further, evaluating
GPU-FRIENDLY KNN AND FOF WITH DUAL TREE WALKS IN JAX/CUDA
interactions between the parent nodes requires only reading each of their children once. This significantly reduces memory access compared to conventional tree walks based on Euler tours where such interactions may be encountered at separate points in time. However, it is worth noting that some threads may be idle if the number of children in par i is smaller than the number of threads in the group. E.g. if we choose a coarsening factor of 8, we’d expect typical nodes to have 8 children which is notably smaller than the minimal number of threads in a group of 32. In principle, this aspect could be further optimized by assigning multiple threads to the same node and then collaboratively inserting neighbours into a joint heap among those threads. However, we do not attempt this optimization here, because it is only a minor concern for leaf(0) leaf interactions (where Nmax ∼ 32 − 64) which tend to dominate the cost of the neighbour search. The C OUNT and I NSERT kernels follow the same structural pattern, but instead of maintaining a heap, they simply count the number of node-node interactions with dlow (node i, node j) ≤ Rmax and insert the corresponding node j indices into the interaction list. Finally, the implementation of the L EAF T O L EAF kernel is again similar to the F IND R MAX kernel. In this case the outer loop runs over query points (assigning one query point per thread) and the inner loop runs over source points. The neighbour heap structure in this case keeps track of kmax radii and point indices that are written out at the end of each query point iteration. To limit register pressure we choose kmax ≤ 32 and call the kernel multiple times if k > kmax , filtering additionally by a minimum radius Rmin (and an equality breaking index offset) that excludes points that were found in previous iterations. C. Multi-GPU Adapting the presented algorithm to multi-GPU scenarios is relatively straightforward. The main idea is that each GPU maintains the local receiving nodes and their corresponding interaction list. Remote source nodes that need to be interacted with are requested once for the evaluation of each plane. Concretely, the main required additions are as follows: (1) We need to additionally store a tuple of two arrays origin = (rank, idx) that saves the origin rank and index for each (unique) source node that appears in the interaction list. (2) When initializing the dense interaction list and super nodes in lines 1-2 of Algorithm 1, Ntopnodes includes all (local or remote) top-nodes and origin must be initialized appropriately. (3) Before line 4 in Algorithm 1 the remote child data n(p) , c(p) , lvl(p) must be requested for each remote origin. The corresponding remote splits spl(p+1) must be communicated as well. The received data is then rearranged such that spl(p+1) correctly indexes contiguous locally available memory. In addition, origin is propagated to the child level. (4) After line 4 in Algorithm 1 all source indices that appear 0 times in the final interaction list can be removed from origin. (5) Before line 6 of Algorithm 1 we need to do a similar request of leaf splits and source point data. The strength of this approach is that remote source nodes required for interactions are requested only once, the number
7
of communication points in the algorithm remains small and predictable and the remaining functions remain exactly identical to the single-GPU case. In the scenarios that we have tested, we find that the additionally required remote data is O(10% - 60%) of the local receiving node data with a notable dependence on the problem setup and the number of source points per GPU (more data tends to imply better balance). Since it is difficult to foresee all the complications that may arise with more complicated setups and at very large GPU counts, we consider the distributed kNN implementation in JZ - TREE to be experimental and preliminary. D. Implementation Details We enhance the presented algorithms with an additional component that allows more efficient early pruning in the iteration through interaction lists. For each interaction we additionally store the interaction radius rlow – corresponding to the lower node-node distance of the interaction. For each receiving node we sort rlow (after line 4 in Algorithm 2) using a bitonic sort network applied to the corresponding segments. We simply initialize these radii to 0 at the top-node level. This improves the performance for two reasons: (1) Since close-by interactions are encountered earlier, the preliminary estimate of Rmax in Algorithm 3 is better and more candidate radii can be discarded early (rather than triggering a more expensive insertion into the neighbour heap). (2) It allows to define an early exit after line 5 of Algorithm 3 and all other kernels that follow a similar structure: If the maximum current estimate of Rmax across all threads is smaller than rlow , we can discard all remaining interactions. In practice, this prunes on the order of 50% of evaluated interactions. Finally, we note again that we need to predict allocation sizes at compile time to enable jit-compilation in JAX. The main additional allocation that we need to predict here is the size of the interaction list source indices isrc (and radii rlow ). Similar to equation (11), we phrase this allocation relative to the estimated node number: N . (17) allocation size ilist = alloc fac ilist · (0) Nmax /2 For d = 3 dimensions, we find that alloc fac ilist ∼ 200 is typically enough, but we note that it is advisable to choose slightly larger values to decrease the chance that the jitcompiled function needs to be aborted due to insufficient available space. Our primary focus in this article is to optimize performance and memory coalescence to point out a path forward to more GPU friendly tree algorithms. However, it is worth noting that the approach at hand does come at a notable memory cost: (0) With Nmax ∼ 48 and alloc fac ilist ∼ 200 the isrc and rlow arrays each require the allocation of about 10 · N integer / floating point numbers. If only a small number of neighbours k ≲ 10 is requested, this may be the peak contribution to the total required allocation. Further, JAX’s memory management system makes it difficult to guarantee that no unnecessary copies of arrays are created. Our implementation in JZ - TREE is therefore relatively memory-intensive – something that we aim to improve in future releases.
GPU-FRIENDLY KNN AND FOF WITH DUAL TREE WALKS IN JAX/CUDA
z-sort node-to-node 6.4 3.5 12.3 tree construction
4 GPU, N = 4 × 107 12 10
36
25.0
leaf-to-leaf 28
104
reorder 2.6
128
Fig. 5. Execution time in ms of different parts of the neighbour search algorithm for returning k = 16 neighbours in input order for a single GPU setup with 107 points (top) and a 4-GPU setup with 4 × 107 points. The expensive final reordering step can be avoided in most applications.
103 Time [ms]
1 GPU, N = 107
8
102
scipy [32CPU] faiss cupy-knn jaxkd-cuda clover jz-tree
101 linear scaling
100 103
104
E. Performance breakdown All performance measurements throughout this article are run on the booster nodes of the Leonardo cluster at CINECA [36]. Each node has a single 32 core Intel Xeon Platinum 8358 processor, four NVIDIA Ampere A100-64 GPUs and 200 Gbps NVIDIA Mellanox HDR InfiniBand connection. Tests with up to 4 GPUs run on a single node, and larger tests run across several nodes (if NGPU ≥ 4). For CPU codes we consider tests for a single core and a 32 core setup on a single node. In Figure 5 we break down the execution time of different steps of a self-neighbour search for a uniform random distribution in three dimensions for a single-GPU and a multiGPU scenario. The single-GPU case highlights the very low cost of the sorting and the tree construction (about 20% of the total). The most expensive part of the algorithm are the leaf-toleaf interactions – comprising approximately 50% of the total execution time. This is expected due to the high computational intensity of this step. However, for the multi-GPU scenario the costs of several steps increases significantly: The z-sort due to the required exchange of points, the tree construction due to the communication step required for regularization and the node-to-node interactions due to multiple required all-to-all communications and the cost of removing unused nodes from the interaction list. Noteworthily, the leaf-to-leaf interactions only require slightly more time, since they only need a single communication with relatively low volume (thanks to efficient pruning from higher levels). The cumulative effect of these steps is an approximate factor 2 decrease in efficiency. However, the most significant increase in execution time is due to the final reordering. This is not too surprising, since bringing the neighbour list into input order requires an extremely high volume communication (recall that these are k = 16 radii and indices per point). Fortunately, in many applications of neighbour search, it is possible to perform a reduction operation while maintaining the neighbour list in z-order and then only communicate back some small summary statistic per point. We provide a simple interface for this recommended usage pattern in JZ - TREE and we output points in z-order for further multi-GPU benchmarks, staying representative of such uses cases.
105
N
106
107
Fig. 6. Comparison to other KNN libraries for a uniform random distribution in d=3 dimensions.
F. Performance comparisons We compare the performance of JZ - TREE for a kNN-search against other publicly available (exact kNN) libraries in Figure 6 for a single GPU setup. The benchmark is to find the k = 30 nearest neighbours4 for a uniform random distribution of points on the range [0, 1]d in d = 3 dimensions for N separate source and query points at float32 precision. In each case we include preparation steps (e.g. sorting and tree building) in the performance measurement, so that this represents fairly the total time that is needed to evaluate one set of source points with one set of query points. However, we exclude the jit compilation time that is necessary in JAX and CUPY implementations. The libraries that we compare to are: (1) SCIPY- CKDTREE – a CPU based kd-tree library implemented as a C++ extension within S CI P Y, operating on N UM P Y arrays [37]. We include measurements for usage of 1 and 32 worker threads. (2) The FAISS library that provides efficient implementations of brute force neighbour search [19], [20]. (3) CUPY- KNN that implements neighbour search through a one-sided traversal of kd-trees in CUDA kernels [13]. (4) Similarly, JAXKD - CUDA based on the CUDA KDT REE library, but offering a convenient jax interface [38]–[40]. (5) CLOVER which traverses a graph based on a random voronoi tessellation [23]. JZ - TREE outperforms all competitor libraries by a significant margin for nearly all problem sizes (except the brute-force approach of FAISS at very small problem sizes N ≲ 104 where the cost of the many kernel launches leads to an irreducible overhead of ∼ 1ms.) For N ≲ 106 CLOVER remains the closest competitor (within about a factor 2), but at larger problem sizes N ≫ 106 CLOVER starts scaling quadratically making it more than an order of magnitude slower at N ∼ 107 . The kd-tree based libraries all exhibit the same (close to linear) asymptotic scaling as JZ - TREE, but with much larger asymptotic constants. The CPU based 4 In general we use k = 16 as a baseline in benchmarks, but here we use k = 30 to allow comparison with the default setup in CLOVER.
grid uniform normal cosm. (no wrapping) cosm. (with wrapping)
102
9
1 GPU 2 GPUs 4 GPUs 8 GPUs
103
Time [ms]
Time [ms]
GPU-FRIENDLY KNN AND FOF WITH DUAL TREE WALKS IN JAX/CUDA
linear scaling 101
16 GPUs 32 GPUs 64 GPUs
102 linear scaling 101
105
106
N
107
108
(a)
105
106
N per GPU
107
108
(b)
Fig. 7. (a) Performance of JZ - TREE for finding k = 16 neighbours of different distributions in d = 3 dimensions. (b) Efficiency scaling for distributed computing for different numbers of GPUs as a function of the number of points per GPU. The method scales well across different problem setups and to large problem sizes on multi-GPU.
SCIPY- CKDTREE turns out more than two orders of magnitude slower than JZ - TREE and the GPU based kd-tree libraries are more than an order of magnitude slower at N ≳ 106 . This improvement may be largely attributed to several key differences in the tree implementation: (1) The much reduced cost of building a tree in a bottom up approach. (2) The reduced algorithmic cost of a dual (versus one-sided) tree walk and (3) the reduced memory access through warp collaborative evaluation and (4) the improved memory coalescence.
G. Performance across domains To demonstrate that the performance benefits are relatively independent of the problem domain, we show in Figure 7(a) performance benchmarks of JZ - TREE for a variety of different setups. In every case we use query points equal to the source points and look for k = 16 neighbours in d = 3 dimensions. The considered scenarios include (1) a uniform grid, (2) the uniform random distribution, (3) a multivariate normal distribution and (4) the final particle distribution from realistic cosmological simulations. The cosmological simulations were run with DISCO-DJ [41] in a Planck (2018) cosmology [42] with a number of particle-mesh cells and the volume of the box chosen proportionally to√the particle count. Specifically we choose the box size as 3 N h−1 Mpc so that the massresolution stays fixed with increasing problem size. For the cosmological simulation we consider two separate scenarios – one where we appropriately include periodic wrapping in the distance calculation of the kNN – and one where we don’t. For all scenarios we have verified the correctness of the returned neighbour lists against SCIPY- CKDTREE. It is evident that JZ - TREE generalizes well over different problem setups with problem-specific performance differences staying well below a factor of two. We note that the most expensive setup – the cosmological simulation with periodic wrapping – owes its 20−30% reduction in efficiency primarily
to the extra-cost in the wrapping calculation (and not so much to the clustering). If we evaluate the same setup without periodic wrapping, the performance is virtually identical to the uniform random distribution at N ≳ 107 . In Figure 7(b) we evaluate the scaling of the multi-GPU implementation of JZ - TREE for a self-query of k = 16 for a uniform random distribution in d = 3 dimensions. In this test we output the 16 output indices and radii per point in z-order. Importantly, the horizontal axis of the plot shows the number of points per GPU so that e.g. the 64 GPU case with N per GPU = 108 evaluates 1.0 · 1011 neighbours (16 neighbours for each of 64 · 108 points) in about 1.3 seconds. The method scales well to a large number of GPUs. The biggest drop in the number of evaluations per GPU per second is seen when going from one to two GPUs leading to an increase in evaluation time at 108 from 585ms up to 928ms – close to a factor of two. This increase comes from the additional algorithmic steps that need to be taken for the distributed computing (like rearranging points, sample sort, adjusting domain boundaries and communication). However, scaling from 2 to 64 GPUs exhibits only an additional decrease in efficiency of 30% (928ms for two GPUs versus 1256ms for 64 at 108 ). For completeness, we provide additional scaling tests with dimension number, neighbour count and query versus source counts in Appendix A. IV. F RIENDS - OF - FRIENDS CLUSTERING As a second example algorithm we describe an efficient implementation of FoF clustering here. The implementation follows very closely the previously outlined dual-tree-traversal pattern plus a well known approach for handling linking relations. We have tested it well in d = 2 and d = 3 dimensions and for periodic and non-periodic boundary conditions, but the implementation should cleanly generalize to higher dimensional setups as well.
GPU-FRIENDLY KNN AND FOF WITH DUAL TREE WALKS IN JAX/CUDA
10
where V is the volume of the simulation box and α is a parameter that is typically chosen to be ∼ 0.2, e.g. [45].
the interaction is discarded – and d ≤ Rlink – where the points’ groups are linked together. After evaluating these interactions the graph is contracted one final time to obtain a unique label for each group. The multi-GPU implementation of the FoF requires some additional effort to distinguish between links that can be resolved locally immediately and those that need to be saved to be resolved globally at a later point (involving communication). However, these details are not very relevant with respect to the focus of this paper, so they will be described in Appendix B.
A. Implementation
B. Catalogue reduction
The connected components of the FoF graph can be conveniently represented through a pointer igroup that is defined per point. If the pointer points to a point itself igroupi = i, we call i ’a root’. Otherwise, it must point to a point that is of the same group and has a lower index. The root of a point’s group can be found by dereferencing the pointer multiple times until it points to itself. All points that have the same root belong to the same group (and vice versa). The FoF implementation follows the same dual tree walk pattern that is outlined in Algorithm 1. However, in addition to the interaction list, the group pointer igroup(p) is carried through the tree-walk and advected from parent to child nodes on every level. It is initialized on the super-node level as a self-pointer. Before the N ODE T O N ODE pass, we perform a PARENT T O N ODE pass that evaluates ( (p+1) spl(p+1) [igroupi [parent(i)]] if linked (p) igroupi = i else (19)
After the group identification, we bring points into group order. That means we perform a stable sort based on the group index, so that the roots of groups remain in z-order with respect to other roots and points in each group form a contiguous block that is internally in z-order. The last group on each device may continue on subsequent devices. Bringing points into group order is useful to make subsequent reduction steps simpler and to make it simple to read the points in different FoF groups separately if the particle data is dumped. Finally, we calculate summary statistics like the total mass, the inertia radius, the center of mass position and the the center of mass velocity (if particle velocities were provided as input). This step can be done almost entirely locally, except for a small communication step related to the last/first group on each task. We provide the option to filter the resulting catalogues by a minimum particle count and choose 20 for this – as is common in the computational cosmology literature – in the following performance tests.
so that for linked nodes it will point to the first child of the root of its parent node. For unlinked nodes it will simply point to the point itself. A root is considered self-linked if it was linked with any other node or if its diagonal extent is smaller than the linking length. The node-to-node interaction distinguishes three cases: (1) If both nodes already point to the same root or if dlow > Rlink , the interaction is discarded. (2) If dup ≤ Rlink , the other node falls fully inside of the linking length and the nodes are linked together. (3) Otherwise the interaction needs to be evaluated at the child level and is inserted into the interaction list. When two nodes are linked together, we first find their roots and then update the higher index root to point towards the lower index root – thereby linking all points in the two groups together. On GPU it is important to protect against data races in this update (between finding the roots and the update, one of the roots may have changed) with atomic compare and swap operations and a repeat on failure. We first launch a kernel to update igroup in this way and afterwards contract the igroup relation in a separate kernel. This is simply done by setting every pointer in igroup to its root. Finally, we count and insert the interactions that need to be checked on the next level. The point-point interactions in the leaf-to-leaf kernel only need to distinguish between two scenarios d > Rlink – where
C. Performance
The goal of a FoF algorithm is to find the connected components of a graph where each point is a node and edges exist between every pair of nodes that is closer than the linking length Rlink [43], [44]. The linking length in cosmological simulations is often chosen relative to the mean separation between points: 1/3 V Rlink = α (18) N
We evaluate the time that is required to obtain the FoF catalogue for the particle distribution from a cosmological simulation (as described in Section III-G). This includes all the necessary steps, i.e. sorting, tree building, the tree walk, the reordering into group order and the final reduction steps. However, we don’t include disk write time in this benchmark. For comparison we test against the single CPU FoF implementation of HFOF [46], the MPI implementation in G AD GET 4 [47] and the single GPU implementation of JFOF [48]. For HFOF we only benchmark the labelling step, since no catalogue reduction is provided – so results are slightly skewed in its favour. For Gadget4, we read in an hdf5 snapshot that we created with DISCODJ and run only the FoF algorithm. Here, we use the timings that are written into stdout, excluding the initial reading of the input, the initial domain decomposition5 and the final writing of the output. We run Gadget once with 1 MPI task and once with 32 MPI tasks on a 32 core node. JFOF is the only other pure GPU FoF code that we are aware of and it is a research-level implementation to enable differentible halo finding [48]. It uses JAX - KD [CUDA] to iteratively link points together by traversing their neighbour 5 We exclude this, since the input is read initially from a single snapshot onto a single task and is very imbalanced through this until after the first domain decomposition.
GPU-FRIENDLY KNN AND FOF WITH DUAL TREE WALKS IN JAX/CUDA
105
Time [ms]
104
883
1283
1803
2523
3603
5123
hfof [1CPU] Gadget4 [1CPU] Gadget4 [32CPU] jfof [1GPU] jz-tree [1GPU]
103 Time [ms]
643
11
103
1GPU 4GPUs 16GPUs 64GPUs
102
linear scaling
102 101
106
N
107
108
(a)
106
107 N per GPU
108
(b)
Fig. 8. Performance of the full FoF algorithm of JZ - TREE - FOF, including catalogue reduction (a) against other libraries and (b) as a function of the number of devices.
graph. The benchmarks required padding with an additional particle to avoid CUDA memory access errors – as suggested by the authors. The resulting measurements are found in Figure 8(a). Similar to the nearest neighbour search, JZ - TREE scales linearly with the problem size once the GPU is fully saturated N ≳ 107 . The performance of JZ - TREE compares favourably with respect to the alternatives. For 5123 the evaluation takes 1.2s which is about 5 times faster than Gadget4 with 32 cores (5.3s), 18 times faster than JFOF (22s), 66 times faster than hfof (82s) and 116 times faster than Gadget 4 with one core (144s). Finally, we show in Figure 8(b) benchmarks for different GPU counts. The efficiency takes the biggest reduction when jumping from 1 node (≤ 4 GPUs) to multiple nodes (> 4 GPUs) where the communication becomes less efficient. The most relevant factor here is probably the increased communication latency in the distributed link insertion and contraction steps. However, the efficiency only decreases in total by a factor 2 − 3 when scaling from 1 to 64 GPUs, allowing us to calculate FoF group catalouges on 20483 points on 64 GPUs in about 3s.
order of magnitude performance improvements over existing GPU codes with great scaling to distributed computation with large numbers of GPUs. The presented algorithms are implemented in the JZ - TREE library, publicly available on GitHub (reference) and PyPI (reference). They can readily be used in HPC simulation schemes that rely on these components like smoothed particle hydrodynamics, self-interacting dark matter simulations and halo finding in cosmological simulations. Finally, we emphasize that JZ - TREE forms a suitable framework for developing efficient GPU implementations of other algorithms that rely on tree representations, such as the fast multipole method which we will discuss in an upcoming publication.
V. C ONCLUSIONS
A PPENDIX A D ETAILED PROFILING OF K NN
Here we have presented a novel approach to construct a plane-based tree hierarchy to enable GPU friendly dual tree walks. Unlike more conventional kd-trees or oct-trees, this tree structure does not partition all of space, has the same depth everywhere and may exhibit a varying number of unequal sized children. It can be constructed in a bottom-up approach with very little additional performance cost after sorted along a Morton z-order curve. The plane hierarchy allows to implement dual tree walks with good thread collaboration and coalescing memory access patterns. We have demonstrated this on two example applications, nearest neighbour search and FoF clustering – yielding
ACKNOWLEDGMENTS This research was funded in whole or in part by the Austrian Science Fund (FWF) [10.55776/ESP705]. We acknowledge access to the EuroHPC supercomputer LEONARDO, hosted by CINECA (Italy) through the AURELIO call. The authors thank Benjamin Horowitz for help with setting up benchmarks for JFOF.
In this appendix we evaluate the performance of the nearest neighbour search in JZ - TREE in dependence on the problem dimension d, the neighbour count k and the number of query points. We show the corresponding tests in Figure 9. Panels (a) and (b) use identical source and query points. Panel (c) varies the number of query points at a fixed number of source points. The scaling with dimension seems to be close to exponential up to d = 6 after which it seems to start saturating. At d = 8 it is still by a factor 10 faster than an evaluation of the same problem with FAISS – which is quite independent of the
GPU-FRIENDLY KNN AND FOF WITH DUAL TREE WALKS IN JAX/CUDA
N = 106
102 faiss jz-tree
101 2
3
4
5 6 dimension
7
8
102 linear scaling
Time [ms]
102
103
Time [ms]
Time [ms]
104
12
Nsrc = 106 Nsrc = 107 xsrc = xquery
101
101 100
(a)
101
k
102
103
105
(b)
106
Nquery
107
(c)
Scaling of the kNN search for 106 points for (a) k = 16 neighbours with dimension (b) d = 3 with neighbour number and (c) k = 16 and d = 3
Fig. 9. and given source count with varying query size.
dimension number. However, that the performance gap to this brute-force approach is only a factor 10 makes it seem likely that per query point about 10% of all source points need to be checked. It is quite possible that the scaling with dimension is notably better for more clustered distributions. However, we note that evaluating high dimensional queries requires a very large allocation for the interaction list, limiting the usefulness of our implementation for d ≫ 3. In panel (b) we show the scaling with neighbour count. At k ≤ 16 the evaluation time is almost independent of the neighbour count. This is likely due to our choice of (0) the leaf-size Nmax = 48, allowing to typically find O(20) neighbours with the same number of leaf-leaf interactions as lower numbers. However, at k ≥ 32 the evaluation cost scales slightly above linear. The asymptotic super-linearity is likely due to the super-imposed effects of the increased number of traversal kernel launches (requiring ⌈k/32⌉ kernel launches) plus the increasing size of the volume that needs to be checked. Finally, in panel (c) we show how the evaluation time depends on the query size for Nsrc = 106 and 107 points. Additionally, we show the performance of a self-query as a black line for reference. For large query sizes Nquery ≫ Nsrc the algorithm takes slightly more time than a self-query with Nquery points. For small query sizes Nquery ≪ Nsrc the performance plateaus at a level similar to the time required for a self-query with Nsrc points. So the performance approximately mirrors the self-query behaviour with max(Nquery , Nsrc ) points – a result of our choice of building the tree based on their joint distribution. For scenarios where a large source distribution needs to be evaluated a large number of times with small query distributions this is clearly not optimal. We may consider offering a different approach for this scenario in future releases. A PPENDIX B M ULTI -GPU F RIENDS - OF -F RIENDS IMPLEMENTATION The distributed FoF implementation uses all of the same adaptations that were described in Section III-C to manage cross-task interactions. However, additional complications arise, because the global root of a node may lie on another
rank and may never have appeared in the interaction list. To address this, we first build a local FoF graph that treats every remote node (or point) initially as a root. Additionally we keep track of a set of edges between pairs of points (ranka , idxa , rankb , idxb ) that represent links that need to be resolved globally later. Whenever the local updates change the label of a node (or point) with remote origin, we store an edge between the rank and index of the first point in the remote node and its new root. After the leaf-leaf-interactions have been evaluated, we perform an additional step that resolves the saved links globally. In this step we replace the local igroup pointer by a label that includes a rank plus an index pointer. Each edge is sent to the larger involved rank. If the pointed location is (still) a root, the edge can be inserted here by updating that label. Under race conditions we simply resolve the lowest proposed update at the same location and consider the other ones unresolved. If the edge could not be inserted here, we update the larger label with the pointed location and we repeat the procedure (sending the edge to the larger involved rank). After all links have been inserted in this way, we contract the global graph. This proceeds by requesting for each unique label that points towards a remote rank the label on that rank and index. If the label is different, the local label is updated and the procedure is repeated for those points until all labels are converged. R EFERENCES [1] K. Cranmer, J. Brehmer, and G. Louppe, “The frontier of simulationbased inference,” Proceedings of the National Academy of Sciences, vol. 117, no. 48, pp. 30 055–30 062, 2020. [Online]. Available: https://www.pnas.org/doi/abs/10.1073/pnas.1912789117 [2] J. Bradbury, R. Frostig, P. Hawkins, M. Johnson, C. Leary, D. Maclaurin, G. Necula, A. Paszke, J. VanderPlas, S. Wanderman-Milne, and Q. Zhang, “Jax: composable transformations of python+numpy programs,” https://github.com/google/jax, 2018. [3] NVIDIA Corporation, “Cuda c++ programming guide,” https://docs. nvidia.com/cuda/, 2024. [4] ——, “Cuda c++ best practices guide,” https://docs.nvidia.com/cuda/ cuda-c-best-practices-guide/, 2024. [5] H. Samet, Foundations of Multidimensional and Metric Data Structures. Morgan Kaufmann, 2006. [6] M. de Berg, O. Cheong, M. van Kreveld, and M. Overmars, Computational Geometry: Algorithms and Applications. Springer, 2008.
GPU-FRIENDLY KNN AND FOF WITH DUAL TREE WALKS IN JAX/CUDA
[7] J. L. Bentley, “Multidimensional binary search trees used for associative searching,” Communications of the ACM, vol. 18, no. 9, pp. 509–517, 1975. [8] J. Barnes and P. Hut, “A hierarchical o(n log n) force-calculation algorithm,” Nature, vol. 324, pp. 446–449, 1986. [9] L. Greengard and V. Rokhlin, “A fast algorithm for particle simulations,” Journal of Computational Physics, vol. 73, pp. 325–348, 1987. [10] T. Karras, “Maximizing parallelism in the construction of bvhs, octrees, and k-d trees,” High Performance Graphics, 2012. [11] C. Lauterbach, M. Garland, S. Sengupta, D. Luebke, and D. Manocha, “Fast bvh construction on gpus,” in Eurographics, 2009. [12] K. Zhou, Q. Hou, R. Wang, and B. Guo, “Real-time kd-tree construction on graphics hardware,” in ACM SIGGRAPH Asia, 2008. [13] J. Jakob and M. Guthe, “Optimizing lbvh-construction and hierarchytraversal to accelerate knn queries on point clouds using the gpu,” in Computer Graphics Forum, vol. 40, no. 1. Wiley Online Library, 2021, pp. 124–137. [14] A. G. Gray and A. W. Moore, “’n-body’ problems in statistical learning,” in Proceedings of the 14th International Conference on Neural Information Processing Systems, ser. NIPS’00. Cambridge, MA, USA: MIT Press, 2000, p. 500–506. [15] R. R. Curtin, W. B. March, P. Ram, D. V. Anderson, A. G. Gray, and C. L. Isbell, “Tree-independent dual-tree algorithms,” in Proceedings of the 30th International Conference on International Conference on Machine Learning - Volume 28, ser. ICML’13. JMLR.org, 2013, p. III–1435–III–1443. [16] R. W. Hockney and J. W. Eastwood, Computer Simulation Using Particles. CRC Press, 1988. [17] M. P. Allen and D. J. Tildesley, Computer Simulation of Liquids. Oxford University Press, 2017. [18] V. Garcia, E. Debreuve, and M. Barlaud, “Fast k nearest neighbor search using gpu,” in IEEE CVPR Workshops, 2008. [19] J. Johnson, M. Douze, and H. Jégou, “Billion-scale similarity search with gpus,” IEEE Transactions on Big Data, vol. 7, no. 3, pp. 535–547, 2021. [20] M. Douze, A. Guzhva, C. Deng, J. Johnson, G. Szilvasy, P.-E. Mazaré, M. Lomeli, L. Hosseini, and H. Jégou, “The faiss library,” 2025. [Online]. Available: https://arxiv.org/abs/2401.08281 [21] H. Jégou, M. Douze, and C. Schmid, “Product quantization for nearest neighbor search,” IEEE Transactions on Pattern Analysis and Machine Intelligence, vol. 33, no. 1, pp. 117–128, 2011. [22] Y. A. Malkov and D. A. Yashunin, “Efficient and robust approximate nearest neighbor search using hierarchical navigable small world graphs,” IEEE Transactions on Pattern Analysis and Machine Intelligence, vol. 42, no. 4, pp. 824–836, 2020. [23] V. Kamel, H. Yan, and S. Chester, “Clover: A gpu-native, spatiograph-based approach to exact knn,” in Proceedings of the 39th ACM International Conference on Supercomputing, ser. ICS ’25. New York, NY, USA: Association for Computing Machinery, 2025, p. 236–249. [Online]. Available: https://doi.org/10.1145/3721145.3730415 [24] G. M. Morton, “A computer oriented geodetic data base and a new technique in file sequencing,” IBM, Tech. Rep., 1966. [25] H. Sagan, Space-Filling Curves. Springer, 1994. [26] M. Bader, Space-Filling Curves: An Introduction with Applications in Scientific Computing. Springer, 2013. [27] H. Samet, Foundations of Multidimensional and Metric Data Structures. Morgan Kaufmann, 2006. [28] M. Connor and P. Kumar, “Fast construction of k-nearest neighbor graphs for point clouds,” IEEE Transactions on Visualization and Computer Graphics, vol. 16, no. 4, pp. 599–608, 2010. [29] “Ieee standard for floating-point arithmetic,” IEEE Std 754-2019 (Revision of IEEE 754-2008), pp. 1–84, 2019. [30] NVIDIA, “Cub: Cuda unbound library,” 2017. [Online]. Available: https://nvlabs.github.io/cub/ [31] H. Shi and J. Schaeffer, “Parallel sorting by regular sampling,” Journal of Parallel and Distributed Computing, vol. 14, no. 4, pp. 361–372, 1992. [Online]. Available: https://www.sciencedirect.com/ science/article/pii/074373159290075X [32] W. D. Frazer and A. C. McKellar, “Samplesort: A sampling approach to minimal storage tree sorting,” J. ACM, vol. 17, no. 3, p. 496–507, Jul. 1970. [Online]. Available: https://doi.org/10.1145/321592.321600 [33] P. Sanders and S. Winkel, “Super scalar sample sort,” in Algorithms – ESA 2004, S. Albers and T. Radzik, Eds. Berlin, Heidelberg: Springer Berlin Heidelberg, 2004, pp. 784–796. [34] M. Mitzenmacher and E. Upfal, Probability and Computing: Randomized Algorithms and Probabilistic Analysis. Cambridge University Press, 2005.
13
[35] P. Ram, D. Lee, W. March, and A. Gray, “Linear-time algorithms for pairwise statistical problems,” in Advances in Neural Information Processing Systems, Y. Bengio, D. Schuurmans, J. Lafferty, C. Williams, and A. Culotta, Eds., vol. 22. Curran Associates, Inc., 2009. [Online]. Available: https://proceedings.neurips.cc/paper files/paper/2009/file/2421fcb1263b9530df88f7f002e78ea5-Paper.pdf [36] M. Turisini, M. Cestari, and G. Amati, “Leonardo: A pan-european pre-exascale supercomputer for hpc and ai applications,” Journal of Large-Scale Research Facilities, vol. 8, p. A186, 2024. [Online]. Available: https://doi.org/10.17815/jlsrf-8-186 [37] P. e. a. Virtanen, “Scipy 1.0: Fundamental algorithms for scientific computing in python,” Nature Methods, vol. 17, pp. 261–272, 2020. [38] B. Dodge, “jaxkd: Minimal JAX implementation of k-nearest neighbors using a k-d tree,” Jul. 2024. [Online]. Available: https: //github.com/dodgebc/jaxkd [39] ——, “jaxkd-cuda: Custom CUDA kernels for JAX k-d tree operations,” 2024, used via the jaxkd interface. [Online]. Available: https://github.com/dodgebc/jaxkd-cuda [40] I. Wald, “A stack-free traversal algorithm for left-balanced k-d trees,” Journal of Computer Graphics Techniques (JCGT), vol. 14, no. 1, pp. 40–54, 2025. [Online]. Available: http://jcgt.org/published/0014/01/03/ [41] F. List, O. Hahn, T. Flöss, and L. Winkler, “DISCO-DJ II: a differentiable particle-mesh code for cosmology,” arXiv e-prints, p. arXiv:2510.05206, Oct. 2025. [42] Planck Collaboration et al., “Planck 2018 results. VI. Cosmological parameters,” A&A, vol. 641, p. A6, Sep. 2020. [43] M. Davis, G. Efstathiou, C. S. Frenk, and S. D. M. White, “The evolution of large-scale structure in a universe dominated by cold dark matter,” The Astrophysical Journal, vol. 292, pp. 371–394, May 1985. [44] J. P. Huchra and M. J. Geller, “Groups of Galaxies. I. Nearby groups,” The Astrophysical Journal, vol. 257, pp. 423–437, Jun. 1982. [45] C. Lacey and S. Cole, “Merger Rates in Hierarchical Models of Galaxy Formation - Part Two - Comparison with N-Body Simulations,” Monthly Notices of the Royal Astronomical Society, vol. 271, p. 676, Dec. 1994. [46] P. Creasey, “Tree-less 3d friends-of-friends using spatial hashing,” Astronomy and Computing, vol. 25, p. 159–167, Oct. 2018. [Online]. Available: http://dx.doi.org/10.1016/j.ascom.2018.09.010 [47] V. Springel, R. Pakmor, O. Zier, and M. Reinecke, “Simulating cosmic structure formation with the GADGET-4 code,” Monthly Notices of the Royal Astronomical Society, vol. 506, no. 2, pp. 2871–2949, Sep. 2021. [48] B. Horowitz and A. E. Bayer, “jFoF: GPU Cluster Finding with Gradient Propagation,” arXiv e-prints, p. arXiv:2510.26851, Oct. 2025.