KnnSearchSYCLImpl.h File Reference

(bba9247 (Tue Jul 28 21:42:19 2026 -0700))

SYCL device kernels for KNN (Direct and AddMM top-k). More...

#include <algorithm>
#include <limits>
#include <oneapi/dpl/algorithm>
#include <oneapi/dpl/execution>
#include <sycl/sycl.hpp>
#include <type_traits>
#include "open3d/core/SYCLContext.h"
#include "open3d/core/nns/NeighborSearchCommon.h"
#include "open3d/utility/Logging.h"

Go to the source code of this file.

Namespaces
namespace	open3d
namespace	open3d::core
namespace	open3d::core::nns

Macros
#define	CALL_UPDATE(Kval)
#define	CALL_FINALIZE(Kval)
#define	CALL_DIRECT(Kval)
#define	CALL_DIM(NDIMVAL)
#define	CALL_SELECT(Kval)

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

SYCL device kernels for KNN (Direct and AddMM top-k).

Included only by KnnSearchOpsSYCL.cpp. Path selection, CUDA comparison, and fixed-radius / hybrid summaries live in the driver file header (KnnSearchOpsSYCL.cpp). Grid search kernels: FixedRadiusSearchSYCLImpl.h.

Kernel groups (called from \ref KnnSearchSYCL)

Path	Key symbols	When (driver)
Direct	DispatchKnnDirect, KnnDirect	dim∈[1,8], k≤32, not

force_addmm_path | | AddMM fused | UpdateTopKFromTile, FinalizeTopK, KBucket | else, k≤512 after CenterPointsAndQueries | | AddMM large-k | SelectTopKQueries, MergeTopKQueries, AddQueryNormsToDistances | else, k>512 |

Direct: sub-group SLM point tiles, per-lane sorted top-K, shuffle-merge; compile-time NDIM and K; no AddMM, no centering.

AddMM fused: oneMKL produces −2qp tiles; UpdateTopKFromTile adds ‖p‖², clamps, updates register or scratch max-heap (FinalizeTopK adds ‖q‖²).

AddMM large-k: per tile SelectTopKQueries then MergeTopKQueries; P8 fallback uses oneDPL partial_sort per query when merge width exceeds mid-k limits.

Distance / top-k conventions (AddMM paths)

Id	Rule
P2	Tile stores −2qp + ‖p‖²; add ‖q‖² once in finalize
C1	Clamp partial distance ≥ 0
C4	Equal distance → smaller global point index wins
C5	Caller passes batch_knn = min(knn, num_points)

Threshold constants: kSYCLKnnSmallKMax (32), kSYCLKnnMidKMax (512).

Macro Definition Documentation

CALL_DIM

#define CALL_DIM

(

NDIMVAL

)

Value:

    DispatchKnnDirectK<T, TIndex, NDIMVAL>(                                    \
            queue, points_ptr, queries_ptr, num_points, num_queries, actual_k, \
            out_dist_ptr, out_idx_ptr, subgroups_per_wg, tile_points)

CALL_DIRECT

#define CALL_DIRECT

(

Kval

)

Value:

    KnnDirect<T, TIndex, NDIM, Kval, SG>(                                      \
            queue, points_ptr, queries_ptr, num_points, num_queries, actual_k, \
            out_dist_ptr, out_idx_ptr, subgroups_per_wg, tile_points)

CALL_FINALIZE

#define CALL_FINALIZE

(

Kval

)

Value:

    FinalizeTopK<T, TIndex, Kval>(queue, num_queries, running_dist_ptr,       \
                                  running_idx_ptr, out_dist_ptr, out_idx_ptr, \
                                  actual_k, query_norms_ptr)

CALL_SELECT

#define CALL_SELECT

(

Kval

)

Value:

    SelectTopKQueriesHeap<T, TIndex, Kval>(                                    \
            queue, distances_ptr, distance_query_stride, num_queries,          \
            num_points, knn, index_offset, out_indices_ptr, out_distances_ptr, \
            out_query_stride, use_threshold, query_norms_ptr, radius_sq,       \
            scalar_threshold)

CALL_UPDATE

#define CALL_UPDATE

(

Kval

)

Value:

    UpdateTopKFromTile<T, TIndex, Kval>(                                      \
            queue, neg2qp_ptr, distance_stride, point_norms_ptr, num_queries, \
            num_points, point_offset, best_dist_ptr, best_idx_ptr,            \
            use_threshold, threshold)