Skip to content

Flat-layout kd-tree to close the normal-estimation gap with open3d #183

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
rajgandhi1 merged 1 commit into from
Jul 7, 2026
Merged

Flat-layout kd-tree to close the normal-estimation gap with open3d #183

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 10 additions & 5 deletions ROADMAP.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,11 +27,16 @@ for the full tables and reproduction command.
These are the concrete, measurable items that move the benchmark and the
credibility story. In rough priority order:

- **Flat-layout kd-tree** — the pointer-based tree is the bottleneck behind the
~1.7x normal-estimation gap. A contiguous, index-based tree (or `kiddo`) is the
single highest-leverage change. → [#176](https://github.com/rajgandhi1/threecrate/issues/176) *(good first issue)*
- **Dense ICP on large clouds** — at parity on small clouds, still ~0.7–0.8x on
KITTI/TUM. Depends partly on the flat kd-tree. → [#177](https://github.com/rajgandhi1/threecrate/issues/177)
- ~~**Flat-layout kd-tree**~~ — **done** ([#176](https://github.com/rajgandhi1/threecrate/issues/176)).
The pointer/`Box` tree is now a contiguous, index-referenced `Vec<KdNode>`. k-NN
results are identical (all 201 algorithm tests pass); a same-machine A/B measured a
consistent **~8–10% speedup on normal estimation and ~5–9% on ICP**. It does **not**
close the Open3D gap on its own — normals are still ~0.5x on large clouds — because
the dominant remaining cost is per-point PCA and single-threaded correspondence
search, not tree layout. That work continues in [#177](https://github.com/rajgandhi1/threecrate/issues/177).
- **Dense ICP on large clouds** — at parity on small clouds, still ~0.7x on
KITTI/TUM even after the flat kd-tree. Parallelising correspondence search and
per-point PCA (rayon) is the next lever. → [#177](https://github.com/rajgandhi1/threecrate/issues/177)
- **Integrate PCL into the benchmark table** — the PCL harness is written and
builds ([`scripts/pcl_bench/`](scripts/pcl_bench)); it just needs to be run in a
shared environment and folded into the published numbers. → [#179](https://github.com/rajgandhi1/threecrate/issues/179)
Expand Down
19 changes: 16 additions & 3 deletions docs/benchmarks.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -114,6 +114,14 @@ benchmark. Each is covered by the existing unit tests (201 passing).
(full-res KITTI 530.8 → 387.3 ms; TUM 1548.6 → 1005.2 ms; nuScenes → parity).
It has negligible effect on normal estimation, which is PCA-bound, not
search-bound — normals are honestly still ~1.7x behind Open3D at full scale.
- **Flat, array-backed kd-tree** (`nearest_neighbor.rs`, [#176]). The tree was
`Box`-pointer-based; it is now a contiguous `Vec<KdNode>` with children
referenced by index, so traversal is cache-friendly. k-NN output is identical
(201 tests pass). A same-machine A/B measured a **consistent ~8–10% speedup on
normals and ~5–9% on ICP** (e.g. normals KITTI 105.9 → 97.3 ms; ICP TUM
900.6 → 828.2 ms). Honestly, this is a real but modest win that does **not**
close the Open3D gap — the remaining cost is per-point PCA and single-threaded
correspondence, tracked in [#177].
- **Outlier removal now uses the KD-tree** instead of brute force
(`filtering.rs`), turning `radius_outlier_removal` and
`statistical_outlier_removal` from O(n²) into O(n log n). Not exercised by the
Expand All @@ -125,9 +133,11 @@ benchmark. Each is covered by the existing unit tests (201 passing).

## Known remaining gaps (honest)

- **Normal estimation is ~1.7x slower than Open3D at full scale.** The hand-rolled
pointer-based kd-tree and per-point PCA are the bottleneck; a flat-layout kd-tree
(or wiring in `kiddo`) is the next step.
- **Normal estimation is still slower than Open3D at full scale** (~0.5x on
TUM/KITTI). The kd-tree is now flat/array-backed (see "What changed"), which
shaved a consistent ~8–10% off but did not close the gap; the dominant remaining
cost is per-point PCA and single-threaded neighbor search, so parallelising those
is the next step ([#177]).
- **Dense ICP still trails Open3D** on large clouds (KITTI/TUM), even after the
k-NN speedup.
- **GPU knn/normals/icp are not competitive yet** (per-call shader/pipeline rebuilds,
Expand Down Expand Up @@ -182,5 +192,8 @@ Swap `--max-points all` for `--max-points 20000` to reproduce the capped table.
- KITTI raw data: https://www.cvlibs.net/datasets/kitti/raw_data.php
- TUM RGB-D download: https://cvg.cit.tum.de/data/datasets/rgbd-dataset/download
- nuScenes paper: https://arxiv.org/abs/1903.11027

[#176]: https://github.com/rajgandhi1/threecrate/issues/176
[#177]: https://github.com/rajgandhi1/threecrate/issues/177
</content>
</invoke>
148 changes: 78 additions & 70 deletions threecrate-algorithms/src/nearest_neighbor.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,31 +4,31 @@ use std::cmp::Ordering;
use std::collections::BinaryHeap;
use threecrate_core::{NearestNeighborSearch, Point3f, Result};

/// KD-Tree node for efficient nearest neighbor search
/// Sentinel used in place of a child index to mean "no child".
const NIL: u32 = u32::MAX;

/// KD-Tree node stored in a flat, contiguous array.
///
/// Children are referenced by index into the owning `Vec<KdNode>` rather than
/// through `Box` pointers. Keeping every node in one allocation makes traversal
/// cache-friendly — neighbour search is the dominant cost in normal estimation
/// and ICP correspondence, so a contiguous layout directly moves those numbers.
#[derive(Debug)]
struct KdNode {
point: Point3f,
original_index: usize, // Store the original index
left: Option<Box<KdNode>>,
right: Option<Box<KdNode>>,
axis: usize, // 0=x, 1=y, 2=z
original_index: usize, // index into the original input slice
left: u32, // child index, or NIL
right: u32, // child index, or NIL
axis: u8, // splitting axis: 0=x, 1=y, 2=z
}

impl KdNode {
fn new(point: Point3f, original_index: usize, axis: usize) -> Self {
Self {
point,
original_index,
left: None,
right: None,
axis,
}
}
}

/// Efficient KD-Tree implementation for nearest neighbor search
/// Efficient KD-Tree implementation for nearest neighbor search.
///
/// Nodes live in a single contiguous `Vec` (`nodes`); children are referenced by
/// index. `root` is the index of the tree root (always `0` when non-empty).
pub struct KdTree {
root: Option<Box<KdNode>>,
nodes: Vec<KdNode>,
root: Option<u32>,
points: Vec<Point3f>, // Keep original points for reference
}

Expand All @@ -37,6 +37,7 @@ impl KdTree {
pub fn new(points: &[Point3f]) -> Result<Self> {
if points.is_empty() {
return Ok(Self {
nodes: Vec::new(),
root: None,
points: Vec::new(),
});
Expand All @@ -48,56 +49,63 @@ impl KdTree {
.map(|(i, &point)| (point, i))
.collect();

let root = Self::build_tree(&mut points_with_indices, 0, 0, points.len() - 1);
let mut nodes: Vec<KdNode> = Vec::with_capacity(points.len());
let root = Self::build_tree(&mut nodes, &mut points_with_indices, 0, 0, points.len() - 1);

Ok(Self {
root: Some(Box::new(root)),
nodes,
root: Some(root),
points: points.to_vec(),
})
}

/// Recursively build the KD-tree
/// Recursively build the KD-tree into the flat `nodes` array, returning the
/// array index of the subtree root spanning `[start, end]`.
///
/// Slots are filled in pre-order, so the overall root lands at index 0.
fn build_tree(
nodes: &mut Vec<KdNode>,
points: &mut [(Point3f, usize)],
depth: usize,
start: usize,
end: usize,
) -> KdNode {
if start == end {
let (point, index) = points[start];
return KdNode::new(point, index, depth % 3);
}

) -> u32 {
let axis = depth % 3;
let median_idx = (start + end) / 2;

// Find the actual median and partition points around it
Self::select_median(points, start, end, median_idx, axis);

let (point, index) = points[median_idx];
let mut node = KdNode::new(point, index, axis);

// Reserve this node's slot before recursing so children can link to it
// (and to each other) by index.
let my_idx = nodes.len() as u32;
nodes.push(KdNode {
point,
original_index: index,
left: NIL,
right: NIL,
axis: axis as u8,
});

// Build left subtree
if median_idx > start {
node.left = Some(Box::new(Self::build_tree(
points,
depth + 1,
start,
median_idx - 1,
)));
}
let left = if median_idx > start {
Self::build_tree(nodes, points, depth + 1, start, median_idx - 1)
} else {
NIL
};

// Build right subtree
if median_idx < end {
node.right = Some(Box::new(Self::build_tree(
points,
depth + 1,
median_idx + 1,
end,
)));
}
let right = if median_idx < end {
Self::build_tree(nodes, points, depth + 1, median_idx + 1, end)
} else {
NIL
};

node
nodes[my_idx as usize].left = left;
nodes[my_idx as usize].right = right;
my_idx
}

/// Select the median element and partition points around it
Expand Down Expand Up @@ -179,13 +187,14 @@ impl NearestNeighborSearch for KdTree {
// distance, so heap ordering and pruning are unaffected. We take the
// square root once per surviving neighbor when building the result.
let mut heap: BinaryHeap<Neighbor> = BinaryHeap::with_capacity(k + 1);
let mut stack: Vec<&KdNode> = Vec::new();
let mut stack: Vec<u32> = Vec::new();

if let Some(ref root) = self.root {
if let Some(root) = self.root {
stack.push(root);
}

while let Some(node) = stack.pop() {
while let Some(idx) = stack.pop() {
let node = &self.nodes[idx as usize];
let dist_sq = Self::distance_squared(&node.point, query);

if heap.len() < k {
Expand All @@ -203,18 +212,18 @@ impl NearestNeighborSearch for KdTree {
}
}

let query_val = query.coords[node.axis];
let node_val = node.point.coords[node.axis];
let query_val = query.coords[node.axis as usize];
let node_val = node.point.coords[node.axis as usize];
let axis_dist = query_val - node_val;
let axis_dist_sq = axis_dist * axis_dist;

// Near child: the half-space the query point lives in.
// Far child: the other half-space, searched only when it could
// contain a point closer than the current k-th nearest.
let (near, far) = if query_val <= node_val {
(&node.left, &node.right)
(node.left, node.right)
} else {
(&node.right, &node.left)
(node.right, node.left)
};

// Push far before near so near is popped first (LIFO), giving the
Expand All @@ -225,13 +234,11 @@ impl NearestNeighborSearch for KdTree {
} else {
true
};
if search_far {
if let Some(ref far_node) = far {
stack.push(far_node);
}
if search_far && far != NIL {
stack.push(far);
}
if let Some(ref near_node) = near {
stack.push(near_node);
if near != NIL {
stack.push(near);
}
}

Expand All @@ -251,37 +258,38 @@ impl NearestNeighborSearch for KdTree {

let radius_sq = radius * radius;
let mut result: Vec<(usize, f32)> = Vec::new();
let mut stack: Vec<&KdNode> = Vec::new();
let mut stack: Vec<u32> = Vec::new();

if let Some(ref root) = self.root {
if let Some(root) = self.root {
stack.push(root);
}

while let Some(node) = stack.pop() {
while let Some(idx) = stack.pop() {
let node = &self.nodes[idx as usize];
let dist_sq = Self::distance_squared(&node.point, query);
if dist_sq <= radius_sq {
result.push((node.original_index, dist_sq.sqrt()));
}

let query_val = query.coords[node.axis];
let node_val = node.point.coords[node.axis];
let query_val = query.coords[node.axis as usize];
let node_val = node.point.coords[node.axis as usize];
let axis_dist = query_val - node_val;

let (near, far) = if query_val <= node_val {
(&node.left, &node.right)
(node.left, node.right)
} else {
(&node.right, &node.left)
(node.right, node.left)
};

// The far subtree can only contain in-radius points when the
// distance to the splitting hyperplane is within the search radius.
if axis_dist * axis_dist <= radius_sq {
if let Some(ref far_node) = far {
stack.push(far_node);
if far != NIL {
stack.push(far);
}
}
if let Some(ref near_node) = near {
stack.push(near_node);
if near != NIL {
stack.push(near);
}
}

Expand Down
Loading