KnnSearchOpsSYCL.cpp File Reference

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

Host driver for SYCL nearest-neighbor search (KNN, fixed-radius, hybrid). More...

#include <algorithm>
#include <functional>
#include <limits>
#include <numeric>
#include <oneapi/dpl/algorithm>
#include <oneapi/dpl/execution>
#include <oneapi/dpl/numeric>
#include <sycl/sycl.hpp>
#include "open3d/core/SYCLUtils.h"
#include "open3d/core/Tensor.h"
#include "open3d/core/linalg/AddMM.h"
#include "open3d/core/nns/FixedRadiusIndex.h"
#include "open3d/core/nns/KnnIndex.h"
#include "open3d/core/nns/NeighborSearchAllocator.h"
#include "open3d/core/nns/kernel/FixedRadiusSearchSYCLImpl.h"
#include "open3d/core/nns/kernel/KnnSearchSYCLImpl.h"
#include "open3d/utility/Logging.h"

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

Macros
#define	INSTANTIATE(T, TIndex)

Functions
template<class T , class TIndex >
void	open3d::core::nns::KnnSearchSYCL (const Tensor &points, const Tensor &points_row_splits, const Tensor &queries, const Tensor &queries_row_splits, int knn, Tensor &neighbors_index, Tensor &neighbors_row_splits, Tensor &neighbors_distance, int64_t tile_bytes, int64_t max_tile_queries, int64_t tile_points_alignment, bool force_addmm_path)
template<class T , class TIndex >
void	open3d::core::nns::FixedRadiusSearchSYCL (const Tensor &points, const Tensor &queries, double radius, const Tensor &points_row_splits, const Tensor &queries_row_splits, const Tensor &hash_table_splits, const Tensor &hash_table_index, const Tensor &hash_table_cell_splits, const Metric metric, const bool ignore_query_point, const bool return_distances, const bool sort, Tensor &neighbors_index, Tensor &neighbors_row_splits, Tensor &neighbors_distance, int64_t)
template<class T , class TIndex >
void	open3d::core::nns::HybridSearchSYCL (const Tensor &points, const Tensor &queries, double radius, int max_knn, const Tensor &points_row_splits, const Tensor &queries_row_splits, const Tensor &hash_table_splits, const Tensor &hash_table_index, const Tensor &hash_table_cell_splits, const Metric metric, Tensor &neighbors_index, Tensor &neighbors_count, Tensor &neighbors_distance, int64_t)
template void	open3d::core::nns::BuildSpatialHashTableSYCL< float > (const Tensor &points, double radius, const Tensor &points_row_splits, const Tensor &hash_table_splits, Tensor &hash_table_index, Tensor &hash_table_cell_splits)
template void	open3d::core::nns::BuildSpatialHashTableSYCL< double > (const Tensor &points, double radius, const Tensor &points_row_splits, const Tensor &hash_table_splits, Tensor &hash_table_index, Tensor &hash_table_cell_splits)

Detailed Description

Host driver for SYCL nearest-neighbor search (KNN, fixed-radius, hybrid).

What lives here

• KnnSearchSYCL — batched k-NN (L2); three internal paths (Direct / AddMM).
• FixedRadiusSearchSYCL — variable-length neighbors within radius (grid).
• HybridSearchSYCL — capped top-max_knn within radius + counts (grid).

Callers: NearestNeighborSearch / KnnIndex / FixedRadiusIndex when dataset and query tensors are on a SYCL device. CPU tensors use NanoFlann.

Includes: KNN device code in kernel/KnnSearchSYCLImpl.h; uniform-grid kernels in kernel/FixedRadiusSearchSYCLImpl.h. Short index: nns/SYCL_DESIGN.md. This file header is the maintainer reference for path selection and end-to-end flow.

Three search modes in this file

Function	Purpose
KnnSearchSYCL	Fixed k nearest neighbors (L2), batched queries ×

dataset | | FixedRadiusSearchSYCL | All neighbors within radius (variable output size per query) | | HybridSearchSYCL | Up to max_knn closest neighbors within radius + in-radius count |

Fixed-radius and hybrid share the same spatial-hash grid index built in FixedRadiusIndex::SetTensorData (see FixedRadiusSearchSYCLImpl.h). The sections below focus on KNN; grid search is summarized at the end.

KNN path selection

For each batch, batch_knn = min(knn, num_points) (convention C5). CUDA KNN uses two strategies (small brute + GEMM + FAISS warp/block select); SYCL uses three paths:

Decision tree

• If force_addmm_path == false and 1 ≤ dim ≤ 8 and batch_knn ≤ 32 → Direct brute-force L2 (UseKnnDirect / DispatchKnnDirect).
• Else center points and queries once (AddMM paths only; see below).
  • If batch_knn ≤ 512 → AddMM fused (running heap per query tile).
  • Else → AddMM large-k (select + merge; may use oneDPL partial_sort).

force_addmm_path=true is for benchmarks only — always skips Direct.

| Path | When | CUDA analogue | SYCL implementation | |---—|---—|------------—|------------------—| | Direct | !force_addmm ∧ dim∈[1,8] ∧ k≤32 | Small dim/k brute kernel | Sub-group SLM tiles + shuffle-merge top-k | | AddMM fused | else, k≤512 | GEMM + FAISS BlockSelect / heap | oneMKL AddMM + fused UpdateTopKFromTile | | AddMM large-k | else, k>512 | Multi-pass FAISS select | Tile select/merge + serial partial_sort (P8) |

Constants: kSYCLKnnSmallKMax = 32, kSYCLKnnMidKMax = 512 (see KnnSearchSYCLImpl.h).

Direct path (<tt>DispatchKnnDirect</tt>)

Idea: Compute L2 as Σ_d (p_d − q_d)² directly — no AddMM, no data centering, no deferred ‖q‖² term.

Algorithm

• One sub-group owns one query; several queries share a work-group.
• Dataset points are loaded in SLM tiles (double-buffered where applicable).
• Each lane maintains a private sorted top-K; sub-group shuffle-merge; lane 0 writes indices and distances.
• Template params: compile-time NDIM (1…8) and K bucket (1…32). Float typically uses sub-group size 16; double prefers 8 when supported.

Typical use: 3D point clouds, small k — often the fastest path on Intel Xe.

AddMM fused path (k ≤ 512)

Idea: Expand squared L2 as ‖q‖² − 2 q·p + ‖p‖². oneMKL AddMM builds the −2 q·p tile for each (query-tile × point-tile) pair; top-k selection fuses ‖p‖² and running heap updates (UpdateTopKFromTile, FinalizeTopK).

Algorithm

1. Center points and queries (subtract first dataset row) — reduces float32 cancellation in ‖p‖² and ‖q‖² when coordinates or features have large norm (CUDA GEMM KNN does not center; Direct SYCL path does not need it).
2. Precompute ‖p‖² per point and ‖q‖² per query (or equivalent norms).
3. For each tile pair: AddMM → partial tile −2qp; add ‖p‖², clamp ≥ 0, update per-query max-heap (size = KBucket(batch_knn)).
4. FinalizeTopK: heap-sort, add ‖q‖², write batch_knn neighbors.

Heap storage: k ≤ 32 → GRF-resident; k ∈ {64…512} → per-work-item scratch.

AddMM large-k path (k > 512)

Idea: Same tiled AddMM and centering as fused path, but SelectTopKQueries

• MergeTopKQueries instead of one fused running heap across all tiles.

Algorithm

• Same centering, norms, and AddMM tile loop.
• Per point-tile: select tile top-k per query; merge into global best-so-far.
• Once: AddQueryNormsToDistances (add ‖q‖², clamp).

P8 caveat: For very large k, merge/select may fall back to per-query serial partial_sort (oneDPL has no cheap segmented top-k) — correct but slow.

KNN: SYCL vs CUDA (design differences)

Topic	CUDA	SYCL (this file)
Small k, low dim	Brute + GEMM paths	+ Direct sub-group path (no

GEMM) | | GEMM | cuBLAS | oneMKL via AddMM | | Top-k on GEMM tiles | FAISS warp/block select | Custom heap / select-merge / oneDPL | | Float32 stability | Uses norms as-is on GEMM path | Centers data before AddMM | | Large k | Multi-pass masking (1024/2048) | Select/merge + partial_sort per query |

AddMM KNN conventions (all non-Direct paths)

Id	Rule
P2	Tile stores partial distance −2qp + ‖p‖²; ‖q‖² added once in

finalize | | C1 | Clamp partial / final distance ≥ 0 | | C4 | Equal distance → smaller point index wins | | C5 | batch_knn = min(knn, num_points) (caller / driver) |

Fixed-radius (\ref FixedRadiusSearchSYCL) — summary

Idea: Return every dataset point within radius of each query; output length varies per query (CSR layout). Ported from CUDA FixedRadiusSearchImpl.cuh; device detail in FixedRadiusSearchSYCLImpl.h.

Grid build (once per index, FixedRadiusIndex::SetTensorData):

• Uniform spatial hash with cell size 2×radius.
• BuildSpatialHashTableSYCL: count points per cell → oneDPL inclusive scan → scatter into CSR (hash_table_cell_splits, hash_table_index).

Query algorithm

1. Count: CountNeighborsSYCL — for each query, visit 8 corner-adjacent bins; count points with squared L2 ≤ radius².
2. Inclusive scan of counts → neighbors_row_splits; allocate index (and optional distance) buffers for the total neighbor count.
3. Gather: WriteNeighborsSYCL — same 8 bins; write indices and squared distances.
4. Optional **sort=true:** SortNeighborsByDistanceSYCL — segmented oneDPL sort_by_key per query (ties not secondarily ordered by index; matches CUDA cub::DeviceSegmentedRadixSort::SortPairs).

Typical use: full neighborhood inside a ball, not a fixed-k cap.

Hybrid (\ref HybridSearchSYCL) — summary

Idea: Count all in-radius neighbors, but store only the closest max_knn in fixed (num_queries, max_knn) tensors plus per-query in-radius counts. Same grid as fixed-radius (BuildSpatialHashTableSYCL).

Algorithm

• Allocate fixed-size index/distance outputs (empty slots use sentinels).
• WriteNeighborsHybridSYCL: one pass over 8 bins per query — running top-max_knn (replace-current-max), full in-radius count, then bounded bubble sort (≤ max_knn) for ascending distance order.

Typical use: FPFH, normals, covariances — “up to K neighbors inside radius.”

Note: Could be expressed as KNN-then-filter; deferred because the grid port already avoids O(queries × points) brute force.

Shared primitive map (CUDA → SYCL)

Role	CUDA	SYCL
Dense GEMM	cuBLAS	oneMKL
Prefix sum	CUB	oneDPL / work-group scan
Segmented sort	CUB radix	oneDPL sort_by_key
Parallel loops	CUDA kernels	SYCL queues / ParallelFor

Maintenance notes

• Direct KNN is the most tuned path; AddMM KNN paths may need further optimization.
• force_addmm_path forces GEMM paths for regression / benchmark parity testing.
• Hybrid could theoretically be “KNN then filter by radius”; deferred — grid port already avoids O(queries × points) brute force for radius-limited search.

Related source files

File	Role
KnnSearchSYCLImpl.h	Direct / AddMM SYCL kernels, KBucket, thresholds

| | FixedRadiusSearchSYCLImpl.h | Grid build, count/write/sort/hybrid kernels | | NearestNeighborSearch.cpp | Device dispatch to this driver | | AddMM.h | oneMKL batched GEMM for KNN tiles |

Macro Definition Documentation

INSTANTIATE

#define INSTANTIATE	(		T,
			TIndex
	)

Value:

    template void KnnSearchSYCL<T, TIndex>(                                    \
            const Tensor& points, const Tensor& points_row_splits,             \
            const Tensor& queries, const Tensor& queries_row_splits, int knn,  \
            Tensor& neighbors_index, Tensor& neighbors_row_splits,             \
            Tensor& neighbors_distance, int64_t tile_bytes,                    \
            int64_t max_tile_queries, int64_t tile_points_alignment,           \
            bool force_addmm_path);                                            \
    template void FixedRadiusSearchSYCL<T, TIndex>(                            \
            const Tensor& points, const Tensor& queries, double radius,        \
            const Tensor& points_row_splits, const Tensor& queries_row_splits, \
            const Tensor&, const Tensor&, const Tensor&, const Metric metric,  \
            const bool ignore_query_point, const bool return_distances,        \
            const bool sort, Tensor& neighbors_index,                          \
            Tensor& neighbors_row_splits, Tensor& neighbors_distance,          \
            int64_t tile_bytes);                                               \
    template void HybridSearchSYCL<T, TIndex>(                                 \
            const Tensor& points, const Tensor& queries, double radius,        \
            int max_knn, const Tensor& points_row_splits,                      \
            const Tensor& queries_row_splits, const Tensor&, const Tensor&,    \
            const Tensor&, const Metric metric, Tensor& neighbors_index,       \
            Tensor& neighbors_count, Tensor& neighbors_distance,               \
            int64_t tile_bytes);