Skip to content

Commit ef3f5a1

Browse files
ssheoreyssheorey
committed
Merge branch 'copilot/port-cuda-kdtree-to-sycl' of github.com:isl-org/Open3D into copilot/add-sycl-kernels-for-cuda
2 parents 3569d35 + 058c09f commit ef3f5a1

13 files changed

Lines changed: 1136 additions & 59 deletions

‎cpp/open3d/core/CMakeLists.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -102,6 +102,7 @@ open3d_sycl_target_sources(core_impl PRIVATE
102102
linalg/LeastSquaresSYCL.cpp
103103
linalg/LUSYCL.cpp
104104
linalg/MatmulSYCL.cpp
105+
nns/KnnSearchOpsSYCL.cpp
105106
linalg/SolveSYCL.cpp
106107
linalg/SVDSYCL.cpp
107108
linalg/TriSYCL.cpp

‎cpp/open3d/core/nns/FixedRadiusIndex.cpp‎

Lines changed: 31 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -104,8 +104,16 @@ bool FixedRadiusIndex::SetTensorData(const Tensor &dataset_points,
104104
CALL_BUILD(double, BuildSpatialHashTableCUDA)
105105
#else
106106
utility::LogError(
107-
"-DBUILD_CUDA_MODULE=OFF. Please compile Open3d with "
107+
"-DBUILD_CUDA_MODULE=OFF. Please compile Open3D with "
108108
"-DBUILD_CUDA_MODULE=ON.");
109+
#endif
110+
} else if (device.IsSYCL()) {
111+
#ifdef BUILD_SYCL_MODULE
112+
return true;
113+
#else
114+
utility::LogError(
115+
"-DBUILD_SYCL_MODULE=OFF. Please compile Open3D with "
116+
"-DBUILD_SYCL_MODULE=ON.");
109117
#endif
110118
} else {
111119
CALL_BUILD(float, BuildSpatialHashTableCPU)
@@ -171,8 +179,18 @@ std::tuple<Tensor, Tensor, Tensor> FixedRadiusIndex::SearchRadius(
171179
});
172180
#else
173181
utility::LogError(
174-
"-DBUILD_CUDA_MODULE=OFF. Please compile Open3d with "
182+
"-DBUILD_CUDA_MODULE=OFF. Please compile Open3D with "
175183
"-DBUILD_CUDA_MODULE=ON.");
184+
#endif
185+
} else if (device.IsSYCL()) {
186+
#ifdef BUILD_SYCL_MODULE
187+
DISPATCH_FLOAT_INT_DTYPE_TO_TEMPLATE(dtype, index_dtype, [&]() {
188+
FixedRadiusSearchSYCL<scalar_t, int_t>(RADIUS_PARAMETERS);
189+
});
190+
#else
191+
utility::LogError(
192+
"-DBUILD_SYCL_MODULE=OFF. Please compile Open3D with "
193+
"-DBUILD_SYCL_MODULE=ON.");
176194
#endif
177195
} else {
178196
DISPATCH_FLOAT_INT_DTYPE_TO_TEMPLATE(dtype, index_dtype, [&]() {
@@ -240,8 +258,18 @@ std::tuple<Tensor, Tensor, Tensor> FixedRadiusIndex::SearchHybrid(
240258
});
241259
#else
242260
utility::LogError(
243-
"-DBUILD_CUDA_MODULE=OFF. Please compile Open3d with "
261+
"-DBUILD_CUDA_MODULE=OFF. Please compile Open3D with "
244262
"-DBUILD_CUDA_MODULE=ON.");
263+
#endif
264+
} else if (device.IsSYCL()) {
265+
#ifdef BUILD_SYCL_MODULE
266+
DISPATCH_FLOAT_INT_DTYPE_TO_TEMPLATE(dtype, index_dtype, [&]() {
267+
HybridSearchSYCL<scalar_t, int_t>(HYBRID_PARAMETERS);
268+
});
269+
#else
270+
utility::LogError(
271+
"-DBUILD_SYCL_MODULE=OFF. Please compile Open3D with "
272+
"-DBUILD_SYCL_MODULE=ON.");
245273
#endif
246274
} else {
247275
DISPATCH_FLOAT_INT_DTYPE_TO_TEMPLATE(dtype, index_dtype, [&]() {

‎cpp/open3d/core/nns/FixedRadiusIndex.h‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -376,6 +376,40 @@ void HybridSearchCUDA(const Tensor& points,
376376
Tensor& neighbors_distance);
377377
#endif
378378

379+
#ifdef BUILD_SYCL_MODULE
380+
template <class T, class TIndex>
381+
void FixedRadiusSearchSYCL(const Tensor& points,
382+
const Tensor& queries,
383+
double radius,
384+
const Tensor& points_row_splits,
385+
const Tensor& queries_row_splits,
386+
const Tensor& hash_table_splits,
387+
const Tensor& hash_table_index,
388+
const Tensor& hash_table_cell_splits,
389+
const Metric metric,
390+
const bool ignore_query_point,
391+
const bool return_distances,
392+
const bool sort,
393+
Tensor& neighbors_index,
394+
Tensor& neighbors_row_splits,
395+
Tensor& neighbors_distance);
396+
397+
template <class T, class TIndex>
398+
void HybridSearchSYCL(const Tensor& points,
399+
const Tensor& queries,
400+
double radius,
401+
int max_knn,
402+
const Tensor& points_row_splits,
403+
const Tensor& queries_row_splits,
404+
const Tensor& hash_table_splits,
405+
const Tensor& hash_table_index,
406+
const Tensor& hash_table_cell_splits,
407+
const Metric metric,
408+
Tensor& neighbors_index,
409+
Tensor& neighbors_count,
410+
Tensor& neighbors_distance);
411+
#endif
412+
379413
/// \class FixedRadiusIndex
380414
///
381415
/// \brief FixedRadiusIndex for nearest neighbor range search.

‎cpp/open3d/core/nns/KnnIndex.cpp‎

Lines changed: 25 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,17 @@ bool KnnIndex::SetTensorData(const Tensor& dataset_points,
6767
utility::LogError(
6868
"GPU Tensor is not supported when -DBUILD_CUDA_MODULE=OFF. "
6969
"Please recompile Open3d With -DBUILD_CUDA_MODULE=ON.");
70+
#endif
71+
} else if (dataset_points.IsSYCL()) {
72+
#ifdef BUILD_SYCL_MODULE
73+
dataset_points_ = dataset_points.Contiguous();
74+
points_row_splits_ = points_row_splits.Contiguous();
75+
index_dtype_ = index_dtype;
76+
return true;
77+
#else
78+
utility::LogError(
79+
"SYCL Tensor is not supported when -DBUILD_SYCL_MODULE=OFF. "
80+
"Please recompile Open3D with -DBUILD_SYCL_MODULE=ON.");
7081
#endif
7182
} else {
7283
utility::LogError(
@@ -124,13 +135,24 @@ std::pair<Tensor, Tensor> KnnIndex::SearchKnn(const Tensor& query_points,
124135
});
125136
#else
126137
utility::LogError(
127-
"-DBUILD_CUDA_MODULE=OFF. Please compile Open3d with "
138+
"-DBUILD_CUDA_MODULE=OFF. Please compile Open3D with "
128139
"-DBUILD_CUDA_MODULE=ON.");
140+
#endif
141+
} else if (device.IsSYCL()) {
142+
#ifdef BUILD_SYCL_MODULE
143+
const Dtype index_dtype = GetIndexDtype();
144+
DISPATCH_FLOAT_INT_DTYPE_TO_TEMPLATE(dtype, index_dtype, [&]() {
145+
KnnSearchSYCL<scalar_t, int_t>(KNN_PARAMETERS);
146+
});
147+
#else
148+
utility::LogError(
149+
"-DBUILD_SYCL_MODULE=OFF. Please compile Open3D with "
150+
"-DBUILD_SYCL_MODULE=ON.");
129151
#endif
130152
} else {
131153
utility::LogError(
132-
"-DBUILD_CUDA_MODULE=OFF. Please compile Open3d with "
133-
"-DBUILD_CUDA_MODULE=ON.");
154+
"KnnIndex only supports CUDA and SYCL tensors. Please use "
155+
"NanoFlannIndex instead for CPU tensors.");
134156
}
135157
return std::make_pair(neighbors_index, neighbors_distance);
136158
}

‎cpp/open3d/core/nns/KnnIndex.h‎

Lines changed: 16 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -29,14 +29,28 @@ void KnnSearchCUDA(const Tensor& points,
2929
Tensor& neighbors_distance);
3030
#endif
3131

32+
#ifdef BUILD_SYCL_MODULE
33+
template <class T, class TIndex>
34+
void KnnSearchSYCL(const Tensor& points,
35+
const Tensor& points_row_splits,
36+
const Tensor& queries,
37+
const Tensor& queries_row_splits,
38+
int knn,
39+
Tensor& neighbors_index,
40+
Tensor& neighbors_row_splits,
41+
Tensor& neighbors_distance);
42+
#endif
43+
3244
class KnnIndex : public NNSIndex {
3345
public:
3446
KnnIndex();
3547

3648
/// \brief Parameterized Constructor.
3749
///
38-
/// \param dataset_points Provides a set of data points as Tensor for KDTree
39-
/// construction.
50+
/// \param dataset_points Provides a set of data points as Tensor for
51+
/// nearest neighbor search. CPU tensors use NanoFlann through
52+
/// open3d::core::nns::NearestNeighborSearch, while CUDA and SYCL tensors
53+
/// are handled by this class.
4054
KnnIndex(const Tensor& dataset_points);
4155
KnnIndex(const Tensor& dataset_points, const Dtype& index_dtype);
4256
~KnnIndex();

0 commit comments

Comments
 (0)