Parallel Batch-Dynamic $k$d-trees

Given a query coordinate $q$, a $k$-NN query finds the $k$ nearest neighbors of $q$ amongst elements of the $k$d-tree by performing a pruned search on the tree [18]. The canonical approach is to traverse the tree while inserting points in the current node into a buffer that maintains only the $k$ nearest neighbors encountered so far. Then, entire subtrees can be pruned during the traversal based on the distance of the $k$’th nearest neighbor found so far.

Besides the canonical approach explained above, there has been a lot of prior work on “dual-tree” approaches [19, 6, 20, 21], which provide speedups on serial batched $k$-NN queries. In particular, the dual-tree approach involves building a second $k$d-tree over the set of query points, and then exploring the two trees simultaneously in order to exploit the spatial partitioning provided by the $k$d-tree structure. We parallelized and tested this dual-tree approach in our work.

We use the following parallel primitives in our algorithms.

• •

  Prefix sum takes as input a sequence of values $[a_{1},a_{2},\ldots,a_{n}]$, an associative binary operator $\oplus$, and an identity $i$, and returns the sequence $[i,a_{1},(a_{1}\oplus a_{2}),...,(a_{1}\oplus a_{2}\oplus\cdots\oplus a_{n-1})]$ as well as the overall sum of the elements. Prefix sum can be implemented in $O(n)$ work and $O(\log{n})$ depth [25].

• •

  Partition takes an array $A$, a predicate function $f$, and a partition value $p$ and outputs a new array such that all the values $a\in A,a<p$ appear before all the values $a\in A,a\geq p$. Partition can be implemented in $O(n)$ work and $O(\log{n})$ depth [25].

• •

  Median partition takes $n$ elements and a comparator and partitions the elements based on the median value in $O(n)$ work and $O(\log n\log\log n)$ depth [25].

The algorithm for parallel construction of the cache-oblivious $k$d-tree is shown in Algorithm 1. The function itself is recursive, and so the top level BuildvEB${}_{\text{S}}$ function allocates space on line 2 and calls the recursive function BuildvEBRecursive${}_{\text{S}}$. Refer to Figure 3 for a graphical representation of this construction.

The recursive function BuildvEBRecursive${}_{\text{S}}$ maintains state with 5 parameters: a point set $Q$, a node index $idx$, a splitting dimension $c$, the number of levels to build $l$, and whether it is building the top or bottom of the tree (indicated by $t$). On line 5, we check for the base case—if the number of levels to build is 1, then we have to construct a node. If this is the top of a tree, then this node will be an internal node, so we perform a parallel median partition in dimension $c$ and save it as an internal node. On the other hand, if this is the bottom of the tree, we construct a leaf node that holds all the points in $Q$. Lines 6–9 form the recursive step. In accordance with the exponential layout [4], we have to first construct the top “half” of the tree and then the bottom “half”. Therefore, on line 6, we compute the number of levels $l_{b}$ in the bottom portion as the hyperceiling¹¹1The hyperceiling of $n$, denoted as $\lceil\lceil n\rceil\rceil$ is the smallest power of 2 that is greater than or equal to $n$, i.e., $2^{\lceil\log{n}\rceil}$. of $\frac{l+1}{2}$ and the remaining number of levels $l_{t}$ in the top portion of the tree as $l-l_{b}$. On line 7, we recursively build the top half of the tree. Then, on line 8, we note that because the top half of the tree is a complete binary tree with $l_{t}$ levels, it will use $2^{l_{t}}-1$ nodes. Therefore, we compute $idx_{b}=idx+2^{l_{t}}-1$, the node index where the bottom half of the tree should start because the trees are laid out consecutively in memory. Finally, on line 9, we construct each of the $2^{l_{t}}$ subtrees that fall under the top half of the tree, each with $l_{b}$ levels. Each of these trees falls into a distinct segment of memory in the array, and so we can perform this construction in parallel across all of the subtrees by precomputing the starting index $idx_{i}$ for each of the $2^{l_{t}}$ subtrees.

We trace this process on an example in Figure 3, in which BuildvEB${}_{\text{S}}$ is called on a set $P$ of 8 points. This spawns a call to BuildvEBRecursive${}_{\text{S}}$($P[0:8]$, 0, 0, 4, bottom). On line 6, we will compute $l_{b}=2$ and $l_{t}=2$, and on line 7, we spawn a recursive call to BuildvEBRecursive${}_{\text{S}}$($P[0:8]$, 0, 0, 2, top). This call is shown as the solid box around the top 3 nodes in Figure 3. In this call, we will hit one further level of recursion before laying out the 3 nodes in indices 0, 1, 2. Then, the original recursive call will proceed to line 8, where it will compute $idx_{b}=3$ as the index to begin laying out the $2^{l_{t}}=4$ bottom subtrees. Finally, on line 9, we will precompute that the starting indices for the 4 bottom subtrees are $(idx_{0},idx_{1},idx_{2},idx_{3})=(3,6,9,12)$. This results in 4 parallel recursive calls, shown in the 4 lower dashed boxes in Figure 3. Each of these recursive calls internally has one more level of recursion to lay out their 3 nodes.

After the vEB-layout $k$d-tree is constructed, it can be queried as a regular $k$d-tree—the only difference is the physical layout of the nodes in memory. The correctness of this recursive algorithm can be seen through induction on the number of levels. In particular, we form two inductive hypotheses:

• •

  BuildvEBRecursive${}_{\text{S}}$($Q$, $idx$, $c$, $l$, top) creates a contiguous, fully-balanced binary tree with $l$ levels rooted at memory location $idx$. Furthermore, this binary tree consists of internal $k$d-tree nodes that equally split the point set $Q$ in half at each level.

• •

  BuildvEBRecursive${}_{\text{S}}$($Q$, $idx$, $c$, $\lceil\log{|Q|}\rceil+1$, bottom) creates a contiguous $k$d-tree with $l$ levels rooted at memory location $idx$.

The base cases, with $l=1$, for these inductive hypotheses are explicitly given on line 5. Then, the inductive step follows easily by noting that the definition of hyperceiling implies that the recursive calls on line 9 are all sized such that $l_{b}=\lceil\log{|Q|_{i}}\rceil+1$.

The algorithm for parallel deletion from a single $k$d-tree is shown in Algorithm 2. The function itself is recursive, so the top level Erase${}_{\text{S}}$ calls the subroutine EraseRecursive${}_{\text{S}}$ on the root node on line 2.

The recursive function EraseRecursive${}_{\text{S}}$ acts on one node at a time, represented by the index $idx$. On line 4, it checks for the base case—if the current node is a leaf node, it simply performs a linear scan to mark any points in the leaf node that are also in $Q$ as deleted. Then, it returns NULL if the entire leaf was emptied; otherwise, it returns the current node $idx$. Lines 5–7 represent the recursive case. First, on line 5, we perform a parallel partition of $Q$ around the current node’s splitting hyperplane. We refer to the lower partition as $Q_{l}$ and the upper partition as $Q_{r}$. On line 6, we recurse on the left and right subtrees in parallel, passing $Q_{l}$ to the left subtree and $Q_{r}$ to the right. Finally, line 7 updates the tree structure. We always ensure that every node has 2 children in order to flatten any unnecessary tree traversal. The return value of EraseRecursive${}_{\text{S}}$ indicates the node that should take the place of $idx$ in the tree (potentially the same node)—a return value of NULL indicates that the entire subtree rooted at $idx$ was removed. So, if both the left and right child are removed, then we can remove the current node as well by returning NULL. On the other hand, if neither the left or right child are removed, then the subtree is still intact, and we simply reset the left and right child pointers of the current node and return the current node $idx$, indicating that it was not removed. Finally, if exactly one of the children was removed, then we remove the current node as well and let the remaining child connect directly to its grandparent—in this way, we remove an unnecessary internal splitting node. We do this by simply returning the non-NULL child, signaling that it will take the place of the current node in the $k$d-tree.

We execute our $k$-NN searches in a data-parallel fashion by parallelizing across all of the query points in a batch. The $k$-NN search for each point is executed serially. We implement a “$k$-NN buffer”, a data structure that maintains a list of the current $k$-nearest neighbors and provide quick insert functionality to test and insert new points if they are closer than the existing set. The data structure maintains an internal buffer of size $2k$. To insert a point, it simply adds that point to the end of the buffer. If the buffer is filled up, then it uses a serial selection algorithm to partition the buffer around the $k$-th nearest element and clears out the remaining $k$ elements. This achieves a serial amortized $O(1)$ runtime (because the selection partition step is $O(k)$ and is only performed for every $k$ insertions).

To implement batched $k$-NN on the $k$d-tree, we perform a $k$-NN search for each individual point in parallel across all the points. We now describe the $k$-NN method (kNN${}_{\text{S}}$) for a single point $p$. We first allocate a $k$-NN buffer for the point. Then, we recursively descend through the $k$d-tree searching for the leaf that $p$ falls into. When we find this leaf, we add all of the points in the leaf to the $k$-NN buffer. Then, as the recursion unfolds, we check whether the $k$-NN buffer has $k$ points. If it does not, we add all the points in the sibling of the current node to the $k$-NN buffer to try to fill up the buffer with nearby points as quickly as possible to improve our estimate of the $k$-th nearest neighbor. Otherwise, we use the current distance of the $k$-th nearest neighbor to prune subtrees in the tree. In particular, if the bounding box of the current subtree is entirely contained within the distance of the $k$-th nearest neighbor, we add all points in the subtree to the $k$-NN buffer. If the bounding box is entirely disjoint, then we prune the subtree. Finally, if they intersect, we recurse on the subtree.

As noted by Bentley [1] and Friedman et al. [18], the work for a single nearest-neighbor query on a $k$d-tree is empirically found to be logarithmic in $n$, so the experimental runtime and scalability are much better than suggested by the worst-case bounds.

Insertions are performed in the style of the logarithmic method [2, 3], with the goal of maintaining the minimum number of full trees within BDL-tree. Thus, upon inserting a batch $B$ of points, we rebuild larger trees if it is possible using the existing points and the newly inserted batch. This is implemented as shown in Algorithm 3, and depicted in Figure 4.

First, on line 2, we build a bitmask $F$ of the current set of full static trees in the logarithmic structure. Then, on line 3, because the buffer $k$d-tree has size $X$, we can add $|P|/X$ to $F$ to compute a new bitmask $F_{new}$ of full trees that would result if we added $|P|$ points to the tree structure. As an implementation detail, note that we first add $|P|\mod{X}$ points to the buffer $k$d-tree—if we fill up the buffer $k$d-tree, then we gather the $X$ points from it and treat them as part of $P$, effectively increasing the size of $P$ by $X$. Then, on line 4, taking the bitwise difference between these two bitmasks gives the set of trees that should be consolidated into new larger trees—specifically, any tree that is set in $F_{new}$ but not in $F$ must be constructed from trees that are set in $F$ but not in $F_{new}$. After determining which trees should be combined into new trees, on line 5 we construct all the new trees in parallel—in parallel for each new tree to be constructed, we deconstruct and gather all the points from trees that are being combined into it and then we construct the new tree over these points and any additional required points from $P$ using Algorithm 1.

Refer to Figure 4 for an example of this insertion method (suppose for this example that $X>2$). In Figure 4(a), the BDL-tree contains $X$ points, giving a bitmask of $F=1$ (because only the smallest tree is in use). If we insert $X+1$ points, then we put one node in the buffer tree and compute $F_{new}=1+\frac{X}{X}=2$, and so we have to deconstruct static tree 0 and build static tree 1, as shown in Figure 4(b). Then, if we insert $X+1$ points again, then we again put one point in the buffer tree and compute $F_{new}=2+\frac{X}{X}=3$, and so we simply construct tree 0 on the $X$ new points (leaving tree 1 intact), as seen in Figure 4(c). Finally, if we then insert $X-1$ points, we note that this would fill the buffer up, so we take 1 point from the buffer and insert $X$ points; then, $F_{new}=3+\frac{X}{X}=4$, and so we deconstruct trees 0, 1 and construct tree 2, as seen in Figure 4(d).

When deleting a batch of points, the goal is to maintain balance within the subtrees. Thus, if any subtree decreases to less than half of its full capacity, we move all the points down to a smaller subtree in order to maintain balance. As seen in Algorithm 4, this is implemented as a three-step process.

On line 2, we call a parallel bulk erase subroutine on each of the individual trees in parallel in order to actually erase the points from the trees. On line 3, we scan the trees in parallel and collect the points from all trees which have been depleted to less than half of their original capacity. Finally, on line 4, we use the Insert${}_{\text{L}}$ routine to reinsert these points into the structure.

In the data-parallel $k$-NN implementation, we parallelize over the set of points given to search for nearest neighbors. First, on line 2, we allocate a $k$-NN buffer for each of the points in $S$. Then, for each of the non-empty trees in BDL-tree, we call the data-parallel $k$-NN subroutine on the individual tree, passing in the set $S$ of points and the set of $k$-NN buffers. Because we reuse the same set of $k$-NN buffers for each underlying $k$-NN call, we eventually end up with the $k$-nearest neighbors across all of the individual trees for each point in $S$.

In this benchmark, we measure the time required to insert 10 batches each containing 10% of the points in the dataset into an initially empty tree for each of our two baselines as well as our BDL-tree. The results using object median and spatial median splitting heuristics are shown in Tables II(a), II(b), respectively. Figure 5(b) shows the scalability of the throughput on the 10M points 7D uniform dataset.

We see that B2 achieves the best performance on batched inserts—this is due to the fact that it does not perform any extra work to maintain balance and simply directly inserts points into the existing spatial structure. BDL achieves the second-best performance—this is due to the fact that it does not have to rebuild the entire tree on every insert, but amortizes the rebuilding work across the batches. Finally, B1 has the worst performance, as it must fully rebuild on every insertion. Similar to construction, we note that the spatial median heuristic performs better in the serial case but has lower scalability. With the object median splitting heuristic, BDL achieves parallel speedup of up to $35.5\times$, with an average speedup of $27.2\times$.

In this benchmark, we measure the performance of our batch insertion implementation as the size of the batch varies from 1M points to 5M points. We repeatedly perform batched inserts of the specified size until the entire dataset has been inserted. We provide plots of the results for the 2D VisualVar dataset in Figure 6(a) and for the 7D Uniform dataset in Figure 6(b). The first striking result is that the throughput decreases for B2 as the batch size increases, while the throughput increases for B1 and BDL. This is due to the fact that the work that B2 performs increases as the batch size grows—it has to do more work to compute spatial splits over larger batches, whereas for small batches it quickly computes the upper spatial splits on the first batch and never recomputes them. As a result, B2 has best throughput at smaller batch sizes, but the worst at large batch sizes. For B1 and BDL, note that the throughput increases as the batch size increases. This is due to the fact that each insert has an associated overhead of recomputing spatial partitions, and so the larger the batch size, the fewer times this overhead is paid. For larger batch sizes, BDL has the best throughput among the three implementations for object median. This is again due to the fact that it amortizes the work across inserts, rather than having to recompute spatial splits over the entire dataset at each insert. Finally, note that in most cases, the spatial median heuristic has a better throughput than its object median counterpart, as it takes less work to compute.

In this benchmark, we measure the time required to delete 10 batches each containing 10% of the points in the dataset from an initially full tree for each of our two baselines as well as the BDL-tree. The results using an object median splitting heuristic are shown in Table III(a) and the results using a spatial median splitting heuristic are shown in Table III(b). Figure 5(c) shows the scalability of the throughput on the 10M point 7D uniform dataset.

We observe that B2 has vastly superior performance—it does almost no work other than tombstoning the deleted points so it is extremely efficient. Next, we see that BDL has the second-best performance, as it amortizes the rebuilding across the batches, rather than having to rebuild across the entire point set for every delete. Finally, B1 has the worst performance as it rebuilds on every delete. With the object median splitting heuristic, BDL achieves parallel speedup of up to $33.1\times$, with an average speedup of $28.5\times$.

In this benchmark, we measure the performance of our batch deletion implementation as the size of the batched update varies from 1M points to 5M points. We provide plots of the results for the 2D VisualVar and 7D Uniform datasets in Figures 6(c), 6(d), respectively. We note again that B2 consistently has the highest throughput, with BDL in second and B1 with the lowest throughput. Furthermore, the throughput of B1 and BDL increases as the batch size increases; this is true for the same reasons as with insertions. For B2, we observe consistent throughput across batch sizes.

In this experiment, we measure the scalability of the $k$-NN operation after constructing each data structure over the entire dataset (in a single batch). The results using an object median splitting heuristic are shown in Table IV(a) and the results using a spatial median splitting heuristic are shown in Table IV(b). Figure 5(d) shows the scalability of the throughput on the 10M point 7D uniform dataset. With the object median heuristic, BDL achieves a parallel speedup of up to $46.1\times$, with an average speedup of $40.0\times$.

The results show that B1 and B2 have similar performance (B2 is slightly faster due to implementation differences). Furthermore, they are both faster than BDL-tree. This is to be expected, because the $k$-NN operation is performed directly over the tree after it is constructed over the entire dataset in a single batch. Thus, both baselines will consist of fully balanced trees and will be able to perform very efficient $k$-NN queries. On the other hand, BDL consists of a set of balanced trees, which adds overhead to the $k$-NN operation, as it must be performed separately on each of these individual trees. However, as the next two benchmarks show, BDL provides superior performance in the case of a mixed set of dynamic batch inserts and deletes interspersed with $k$-NN queries.

In this experiment, we benchmark the throughput of the $k$-NN operation on 36 cores with hyper-threading as $k$ varies from 2 to 11. For all three trees, we perform the $k$-NN operation after building the tree from a set of batch insertions, with a batch size of 5% of the dataset, until the entire dataset is inserted. The results are shown for the 2D VisualVar dataset in Figure 7(a) and for the 7D Uniform dataset in Figure 7(b). In both scenarios, we see that B1 has the best $k$-NN performance, followed closely by BDL. B2 has significantly worse performance—this is because the construction of the tree was performed with a set of batch insertions, rather than a single construction over the entire dataset, the tree ends up imbalanced and the $k$-NN query performance suffers.

In this final experiment, we measure the $k$-NN and overall performance of the trees as a mixed set of batch insertions, batch deletions, and batch $k$-NN queries are performed. In particular, we perform a set of 20 batch insertions, each consisting of 5% of the dataset, into the tree. After every 5 batch insertions, we perform a $k$-NN query with $k=5$ (over the entire dataset) to measure the current query performance. Then, we perform a set of 15 batch deletes, each consisting of a random 5% of the dataset (with no repeats). After every 5 batch deletes, we again perform a $k$-NN ($k=5$) query. Overall, there are 7 $k$-NN queries. We thus split the experiment into 7 distinct sections, demarcated by the $k$-NN queries. For each section (labeled as INS0, INS1, INS2, INS3, DEL0, DEL1, DEL2), we measure the $k$-NN runtime as well as the time for the 5 batched insertion/deletion operations. The results are shown for the 2D VisualVar dataset in Figure 8(a) and for the 5d Uniform dataset in Figure 8(b) (we observed similar results for other datasets). The $x$-axis shows the 7 sections, and the $y$-axis shows the time for each of these sections. There are two lines for each tree—a dashed line indicating just the $k$-NN times at each section, and a solid line indicating the total time for the batched update and $k$-NN.

In the Uniform case, we see that B2 performs the worst overall, due almost entirely to its poor $k$-NN performance after batched updates. Note that the batched updates themselves contribute minimally to the total runtime of this baseline—they are very fast but cause significant imbalance in the tree structure, leading to degraded query performance. B1 has the best raw $k$-NN query time, but its overall runtime is the second worst, as the batched updates are quite expensive (in order to maintain the balance that results in fast query times). Finally, BDL represents the best tradeoff between dynamic batch updates and $k$-NN performance. In particular, it has the best total runtime after every operation. The results are quite similar for the VisualVar case, except that we note that B1 has the overall worst runtime, due to very expensive updates.

In addition to the data-parallel approach discussed above, we also implemented $k$-NN based on the dual-tree traversal proposed by March et al. [21]. While we found our dual-tree $k$-NN for BDL was faster on a single-thread due to more efficient pruning, it was not as scalable as other $k$-NN approaches and thus did not perform as well in parallel. Therefore, we omit the experimental results for dual-tree $k$-NN.