Tile sizing (AddMM path)

Data Structures
class	open3d::core::nns::KnnDirectKernel< T, TIndex, NDIM, K, SG >
	Named kernel tag for KnnDirect (SYCL kernel naming). More...

Functions
void	open3d::core::nns::ChooseTileSize (int64_t num_queries, int64_t num_points, int64_t element_size, int64_t tile_bytes, int64_t &tile_queries, int64_t &tile_points, int64_t max_tile_queries=128, int64_t tile_points_alignment=128)
template<typename T , typename TIndex , int K>
void	open3d::core::nns::HeapifyDown (T *d, TIndex *idx, int root)
template<typename T , typename TIndex , int K>
void	open3d::core::nns::HeapSort (T *d, TIndex *idx)
	Heap-sort a compile-time max-heap of size K into ascending order.
template<typename T , typename TIndex , int K>
void	open3d::core::nns::UpdateTopKFromTile (sycl::queue &queue, const T *neg2qp_ptr, int64_t distance_stride, const T *point_norms_ptr, int64_t num_queries, int64_t num_points, TIndex point_offset, T *best_dist_ptr, TIndex *best_idx_ptr, bool use_threshold, T threshold)
template<typename T , typename TIndex , int K>
void	open3d::core::nns::FinalizeTopK (sycl::queue &queue, int64_t num_queries, const T *running_dist_ptr, const TIndex *running_idx_ptr, T *out_dist_ptr, TIndex *out_idx_ptr, int64_t actual_k, const T *query_norms_ptr)
int64_t	open3d::core::nns::KBucket (int64_t k)
	Return the smallest dispatch-bucket value ≥ k.
template<typename T , typename TIndex >
void	open3d::core::nns::DispatchUpdateTopKFromTile (sycl::queue &queue, const T *neg2qp_ptr, int64_t distance_stride, const T *point_norms_ptr, int64_t num_queries, int64_t num_points, int64_t k_bucket, TIndex point_offset, T *best_dist_ptr, TIndex *best_idx_ptr, bool use_threshold, T threshold)
	Instantiate UpdateTopKFromTile for the given k_bucket.
template<typename T , typename TIndex >
void	open3d::core::nns::DispatchFinalizeTopK (sycl::queue &queue, int64_t num_queries, const T *running_dist_ptr, const TIndex *running_idx_ptr, T *out_dist_ptr, TIndex *out_idx_ptr, int64_t actual_k, int64_t k_bucket, const T *query_norms_ptr)
	Instantiate FinalizeTopK for the given k_bucket.
template<typename T , typename TIndex , int NDIM, int K, int SG>
void	open3d::core::nns::KnnDirect (sycl::queue &queue, const T *points_ptr, const T *queries_ptr, int64_t num_points, int64_t num_queries, int64_t actual_k, T *out_dist_ptr, TIndex *out_idx_ptr, int64_t subgroups_per_wg, int64_t tile_points)
	Launch direct-distance KNN for fixed compile-time NDIM, K, and SG.
template<typename T , typename TIndex , int NDIM, int SG>
void	open3d::core::nns::DispatchKnnDirectKForSG (sycl::queue &queue, const T *points_ptr, const T *queries_ptr, int64_t num_points, int64_t num_queries, int64_t actual_k, T *out_dist_ptr, TIndex *out_idx_ptr, int64_t subgroups_per_wg, int64_t tile_points)
template<typename T , typename TIndex , int NDIM>
void	open3d::core::nns::DispatchKnnDirectK (sycl::queue &queue, const T *points_ptr, const T *queries_ptr, int64_t num_points, int64_t num_queries, int64_t actual_k, T *out_dist_ptr, TIndex *out_idx_ptr, int64_t subgroups_per_wg, int64_t tile_points)
template<typename T , typename TIndex >
void	open3d::core::nns::DispatchKnnDirect (sycl::queue &queue, const T *points_ptr, const T *queries_ptr, int64_t dim, int64_t num_points, int64_t num_queries, int64_t actual_k, T *out_dist_ptr, TIndex *out_idx_ptr, int64_t subgroups_per_wg=kKnnDirectSubgroupsPerWG, int64_t tile_points=kKnnDirectTilePoints)
bool	open3d::core::nns::UseKnnDirect (int64_t dim, int64_t knn)
	True if (dim, knn) qualifies for the direct-distance SYCL KNN path.
template<typename T , typename TIndex >
void	open3d::core::nns::DispatchSelectTopKQueries (sycl::queue &queue, const T *distances_ptr, int64_t distance_query_stride, int64_t num_queries, int64_t num_points, int64_t knn, int64_t k_bucket, TIndex index_offset, TIndex *out_indices_ptr, T *out_distances_ptr, int64_t out_query_stride, bool use_threshold, const T *query_norms_ptr, T radius_sq, T scalar_threshold)
template<typename T , typename TIndex >
void	open3d::core::nns::SelectTopKQueries (const Device &device, const T *distances_ptr, int64_t distance_query_stride, int64_t num_queries, int64_t num_points, int64_t knn, TIndex index_offset, TIndex *scratch_indices_ptr, int64_t scratch_query_stride, TIndex *out_indices_ptr, T *out_distances_ptr, int64_t out_query_stride, bool use_threshold=false, const T *query_norms_ptr=nullptr, T radius_sq=T(0), T scalar_threshold=T(0))
template<typename T , typename TIndex >
void	open3d::core::nns::MergeTopKQueries (const Device &device, const T *curr_dist_ptr, const TIndex *curr_idx_ptr, int64_t curr_stride, const T *cand_dist_ptr, const TIndex *cand_idx_ptr, int64_t cand_stride, int64_t num_queries, int64_t knn, TIndex *scratch_ptr, int64_t scratch_stride, TIndex *out_idx_ptr, T *out_dist_ptr, int64_t out_stride)
template<typename T , typename TIndex >
void	open3d::core::nns::AddQueryNormsToDistances (const Device &device, int64_t num_queries, int64_t knn, const TIndex *indices_ptr, T *distances_ptr, const T *query_norms_ptr)

Variables
constexpr int64_t	open3d::core::nns::kKnnDirectSubgroupSize = 16
	Default sub-group width for the direct KNN kernel (float path).
constexpr int64_t	open3d::core::nns::kKnnDirectSubgroupsPerWG = 32
	Default sub-groups per work-group (512 work-items at SG=16).
constexpr int64_t	open3d::core::nns::kKnnDirectTilePoints = 2048
	Default point tile size for SLM staging.
constexpr int64_t	open3d::core::nns::kKnnDirectMaxDim = 8
	Maximum point dimension compiled for DispatchKnnDirect.

Detailed Description

C4 tie-break: smaller global index wins on equal distance.

Dispatch buckets: 1, 2, 4, …, 512 (round knn up with KBucket first).

Cooperative SLM-tiled brute-force KNN for low dimension and low k (D ≤ kKnnDirectMaxDim, K ≤ kSYCLKnnSmallKMax). Accumulates |p−q|² directly (no |q|²−2qp+|p|² cancellation); centering and P2 deferral are not required on this path.

One sub-group of width SG handles one query; subgroups_per_wg queries share a work-group and cooperatively cache each point tile in SLM. Double-buffered tiles hide load latency. Per-lane top-K arrays merge via sub-group shuffle (select_from_group); lane 0 writes the result.

Tuned defaults for dim=3 on Intel Xe: kKnnDirectSubgroupsPerWG and kKnnDirectTilePoints.

Mid/large-k paths apply C1 clamp, C4 tie-break, and P2 (no |q|² in tiles).

Function Documentation

AddQueryNormsToDistances()

template<typename T , typename TIndex >

void open3d::core::nns::AddQueryNormsToDistances	(	const Device &	device,
		int64_t	num_queries,
		int64_t	knn,
		const TIndex *	indices_ptr,
		T *	distances_ptr,
		const T *	query_norms_ptr
	)

Add |q|² to partial distances and clamp ≥ 0 (C1). Called once after all point tiles for the SelectTopK / Merge path (mid-K and large-K).

ChooseTileSize()

void open3d::core::nns::ChooseTileSize ( int64_t  num_queries, int64_t  num_points, int64_t  element_size, int64_t  tile_bytes, int64_t &  tile_queries, int64_t &  tile_points, int64_t  max_tile_queries = 128, int64_t  tile_points_alignment = 128  )

inline

Compute tile dimensions that bound the −2*q*p tile to tile_bytes. tile_queries is capped at max_tile_queries (a good oneMKL GEMM row-tile width; 128 by default). tile_points is rounded down to a multiple of tile_points_alignment (128 by default) once it exceeds that alignment, keeping the column tile GEMM/cache-line friendly.

DispatchFinalizeTopK()

template<typename T , typename TIndex >

void open3d::core::nns::DispatchFinalizeTopK	(	sycl::queue &	queue,
		int64_t	num_queries,
		const T *	running_dist_ptr,
		const TIndex *	running_idx_ptr,
		T *	out_dist_ptr,
		TIndex *	out_idx_ptr,
		int64_t	actual_k,
		int64_t	k_bucket,
		const T *	query_norms_ptr
	)

Instantiate FinalizeTopK for the given k_bucket.

DispatchKnnDirect()

template<typename T , typename TIndex >

void open3d::core::nns::DispatchKnnDirect	(	sycl::queue &	queue,
		const T *	points_ptr,
		const T *	queries_ptr,
		int64_t	dim,
		int64_t	num_points,
		int64_t	num_queries,
		int64_t	actual_k,
		T *	out_dist_ptr,
		TIndex *	out_idx_ptr,
		int64_t	subgroups_per_wg = kKnnDirectSubgroupsPerWG,
		int64_t	tile_points = kKnnDirectTilePoints
	)

Dispatch the direct-distance KNN path by compile-time dimension (1..8) and K-bucket (≤ 32). Writes directly into the caller-allocated output buffers; no public API / build plumbing changes are required. The double-precision sub-group width (8 vs 16) is chosen at runtime inside DispatchKnnDirectK by querying the device; float always uses width 16.

DispatchKnnDirectK()

template<typename T , typename TIndex , int NDIM>

void open3d::core::nns::DispatchKnnDirectK	(	sycl::queue &	queue,
		const T *	points_ptr,
		const T *	queries_ptr,
		int64_t	num_points,
		int64_t	num_queries,
		int64_t	actual_k,
		T *	out_dist_ptr,
		TIndex *	out_idx_ptr,
		int64_t	subgroups_per_wg,
		int64_t	tile_points
	)

Choose sub-group width for direct KNN: float uses 16; double uses 8 when supported, else 16 (see file comment in implementation).

DispatchKnnDirectKForSG()

template<typename T , typename TIndex , int NDIM, int SG>

void open3d::core::nns::DispatchKnnDirectKForSG	(	sycl::queue &	queue,
		const T *	points_ptr,
		const T *	queries_ptr,
		int64_t	num_points,
		int64_t	num_queries,
		int64_t	actual_k,
		T *	out_dist_ptr,
		TIndex *	out_idx_ptr,
		int64_t	subgroups_per_wg,
		int64_t	tile_points
	)

K-bucket dispatch for a fixed compile-time sub-group width SG. Kept separate from DispatchKnnDirectK so the same K-bucket switch can be instantiated at two different SG widths (see DispatchKnnDirectK) without duplicating the bucket logic.

DispatchSelectTopKQueries()

template<typename T , typename TIndex >

void open3d::core::nns::DispatchSelectTopKQueries	(	sycl::queue &	queue,
		const T *	distances_ptr,
		int64_t	distance_query_stride,
		int64_t	num_queries,
		int64_t	num_points,
		int64_t	knn,
		int64_t	k_bucket,
		TIndex	index_offset,
		TIndex *	out_indices_ptr,
		T *	out_distances_ptr,
		int64_t	out_query_stride,
		bool	use_threshold,
		const T *	query_norms_ptr,
		T	radius_sq,
		T	scalar_threshold
	)

K-bucket dispatch to SelectTopKQueriesHeap (file-local, anonymous namespace).

DispatchUpdateTopKFromTile()

template<typename T , typename TIndex >

void open3d::core::nns::DispatchUpdateTopKFromTile	(	sycl::queue &	queue,
		const T *	neg2qp_ptr,
		int64_t	distance_stride,
		const T *	point_norms_ptr,
		int64_t	num_queries,
		int64_t	num_points,
		int64_t	k_bucket,
		TIndex	point_offset,
		T *	best_dist_ptr,
		TIndex *	best_idx_ptr,
		bool	use_threshold,
		T	threshold
	)

Instantiate UpdateTopKFromTile for the given k_bucket.

FinalizeTopK()

template<typename T , typename TIndex , int K>

void open3d::core::nns::FinalizeTopK	(	sycl::queue &	queue,
		int64_t	num_queries,
		const T *	running_dist_ptr,
		const TIndex *	running_idx_ptr,
		T *	out_dist_ptr,
		TIndex *	out_idx_ptr,
		int64_t	actual_k,
		const T *	query_norms_ptr
	)

Heap-sort the running top-K, add |q|² (P2), clamp ≥ 0 (C1), and write the first actual_k entries to the output buffers. Called once after all point tiles (small-k KNN path only).

Parameters

running_dist_ptr	(num_q, K) – max-heap from UpdateTopKFromTile.
running_idx_ptr	(num_q, K) – corresponding global indices.
out_dist_ptr	(num_q, actual_k) – final sorted distances.
out_idx_ptr	(num_q, actual_k) – final sorted indices.
actual_k	Real knn value (≤ K).
query_norms_ptr	|q|² per query (nullptr to skip, e.g. threshold).

HeapifyDown()

template<typename T , typename TIndex , int K>

void open3d::core::nns::HeapifyDown ( T *  d, TIndex *  idx, int  root  )

inline

Heapify-down for a compile-time max-heap of size K. K is a constant so the compiler can unroll the loop and keep d[]/idx[] in GRF for K ≤ 32.

HeapSort()

template<typename T , typename TIndex , int K>

void open3d::core::nns::HeapSort ( T *  d, TIndex *  idx  )

inline

Heap-sort a compile-time max-heap of size K into ascending order.

KBucket()

int64_t open3d::core::nns::KBucket ( int64_t  k )

inline

Return the smallest dispatch-bucket value ≥ k.

KnnDirect()

template<typename T , typename TIndex , int NDIM, int K, int SG>

void open3d::core::nns::KnnDirect	(	sycl::queue &	queue,
		const T *	points_ptr,
		const T *	queries_ptr,
		int64_t	num_points,
		int64_t	num_queries,
		int64_t	actual_k,
		T *	out_dist_ptr,
		TIndex *	out_idx_ptr,
		int64_t	subgroups_per_wg,
		int64_t	tile_points
	)

Launch direct-distance KNN for fixed compile-time NDIM, K, and SG.

MergeTopKQueries()

template<typename T , typename TIndex >

void open3d::core::nns::MergeTopKQueries	(	const Device &	device,
		const T *	curr_dist_ptr,
		const TIndex *	curr_idx_ptr,
		int64_t	curr_stride,
		const T *	cand_dist_ptr,
		const TIndex *	cand_idx_ptr,
		int64_t	cand_stride,
		int64_t	num_queries,
		int64_t	knn,
		TIndex *	scratch_ptr,
		int64_t	scratch_stride,
		TIndex *	out_idx_ptr,
		T *	out_dist_ptr,
		int64_t	out_stride
	)

Merge two sorted (ascending) per-query top-K arrays. Uses a linear merge for knn ≤ kSYCLKnnMidKMax, else an oneDPL partial_sort fallback (P8).

SelectTopKQueries()

template<typename T , typename TIndex >

void open3d::core::nns::SelectTopKQueries	(	const Device &	device,
		const T *	distances_ptr,
		int64_t	distance_query_stride,
		int64_t	num_queries,
		int64_t	num_points,
		int64_t	knn,
		TIndex	index_offset,
		TIndex *	scratch_indices_ptr,
		int64_t	scratch_query_stride,
		TIndex *	out_indices_ptr,
		T *	out_distances_ptr,
		int64_t	out_query_stride,
		bool	use_threshold = false,
		const T *	query_norms_ptr = nullptr,
		T	radius_sq = T(0),
		T	scalar_threshold = T(0)
	)

Select the smallest knn partial distances per query. P2: distances_ptr contains partial dists (no |q|²); callers add it after. C1: clamp each distance to max(0, ...) before comparing. C4: equal distances resolved by global index.

Parameters

query_norms_ptr

Per-query |q|² (size num_queries). When use_threshold is true and this is non-null, threshold is radius_sq − query_norms_ptr[q] (P2 hybrid). Otherwise scalar_threshold.

UpdateTopKFromTile()

template<typename T , typename TIndex , int K>

void open3d::core::nns::UpdateTopKFromTile	(	sycl::queue &	queue,
		const T *	neg2qp_ptr,
		int64_t	distance_stride,
		const T *	point_norms_ptr,
		int64_t	num_queries,
		int64_t	num_points,
		TIndex	point_offset,
		T *	best_dist_ptr,
		TIndex *	best_idx_ptr,
		bool	use_threshold,
		T	threshold
	)

Update the per-query running top-K heap using one point-column tile.

One work-item per query; K is the compile-time heap capacity (= dispatch bucket ≥ actual knn).

K ≤ kSYCLKnnSmallKMax (32): d[K]/idx[K] live in GRF – no scratch spill. K ≤ kSYCLKnnMidKMax (512): spill to scratch, proportional to K.

Fuses three passes into one: Add_(point_norms) + SelectTopKQueries + MergeTopKQueries.

P2: |q|² NOT added; callers add it once in FinalizeTopK. C1: partial_dist = max(0, −2qp + |p|²). C4: equal distances resolved by smaller global point index.

Parameters

neg2qp_ptr	−2*q*p tile from AddMM, shape (num_q, dist_stride).
point_norms_ptr	|p|² for this tile, shape (num_points,).
point_offset	Global index of this tile's first point.
best_dist_ptr	Running top-K distances (max-heap), (num_q, K).
best_idx_ptr	Running top-K global indices, (num_q, K).
use_threshold	True for radius / hybrid search.
threshold	Adjusted per-query threshold = radius² − |q|². Caller passes the per-query array value for each work-item's own query; the scalar path is used by the hybrid SelectTopKQueries variant.

UseKnnDirect()

bool open3d::core::nns::UseKnnDirect ( int64_t  dim, int64_t  knn  )

inline

True if (dim, knn) qualifies for the direct-distance SYCL KNN path.

Variable Documentation

kKnnDirectMaxDim

constexpr int64_t open3d::core::nns::kKnnDirectMaxDim = 8

constexpr

Maximum point dimension compiled for DispatchKnnDirect.

kKnnDirectSubgroupSize

constexpr int64_t open3d::core::nns::kKnnDirectSubgroupSize = 16

constexpr

Default sub-group width for the direct KNN kernel (float path).

kKnnDirectSubgroupsPerWG

constexpr int64_t open3d::core::nns::kKnnDirectSubgroupsPerWG = 32

constexpr

Default sub-groups per work-group (512 work-items at SG=16).

kKnnDirectTilePoints

constexpr int64_t open3d::core::nns::kKnnDirectTilePoints = 2048

constexpr

Default point tile size for SLM staging.