GPU geometry for PyTorch point clouds. Fast k-NN, farthest-point sampling, interpolation, and ray tracing.

Documentation · API reference · Benchmarks

Performance

20x faster k-NN than torch_cluster GPU and 9.2x faster approximate FPS than fpsample CPU, averaged across all 20 benchmark workloads.

About

torchbvh provides CUDA operations for 2D and 3D float32 point clouds, including fixed-size batches:

• k-NN: BVH-accelerated nearest-neighbor search.
• FPS: exact and bucketed approximate farthest-point sampling.
• Interpolation: moving least squares (MLS) over scattered point features.
• Ray tracing: segment and triangle intersections.

Installation

Install CUDA-enabled PyTorch, a matching CUDA toolkit with NVCC, and a supported C++ compiler, then install torchbvh:

pip install torchbvh

See the build guide for compiler and source-install details.

Quickstart

import torch
import torchbvh as tb

points = torch.rand(10_000, 3, device="cuda")
queries = torch.rand_like(points)

with tb.BVH(points) as bvh:
    neighbors, distances_sq = bvh.knn(queries, k=4)

sample = tb.fps(points, target_tokens=2_500, mode="approx_bucketed")
features = torch.rand(10_000, 64, device="cuda")
values = tb.mls_interpolate(points, queries, features, k=4)

See the examples for exact FPS, ray tracing, and gradients.

References

The BVH implementation builds on Chitalu, Dubach, and Komura (2020) and ImplicitBVH.jl. Released under the MIT license.