Splitting fused l2 knn specializations

README.md

@@ -63,6 +63,6 @@ The folder structure mirrors other RAPIDS repos (cuDF, cuML, cuGraph...), with t
 
 - `ci`: Scripts for running CI in PRs
 - `conda`: conda recipes and development conda environments
-- `cpp`: Source code for all C++ code. The code is currently header-only, therefore it is in the `include` folder (with no `src`).
+- `cpp`: Source code for all C++ code. Headers are in the `include` folder and compiled template specializations for the shared libraries are in `src`.
 - `docs`: Source code and scripts for building library documentation
 - `python`: Source code for all Python source code.

cpp/CMakeLists.txt

@@ -203,7 +203,10 @@ if(RAFT_COMPILE_LIBRARIES)
   add_library(raft_nn_lib SHARED
     src/nn/specializations/ball_cover.cu
     src/nn/specializations/detail/ball_cover_lowdim.cu
-    src/nn/specializations/fused_l2_knn.cu
+    src/nn/specializations/fused_l2_knn_long_float_true.cu
+    src/nn/specializations/fused_l2_knn_long_float_false.cu
+    src/nn/specializations/fused_l2_knn_int_float_true.cu
+    src/nn/specializations/fused_l2_knn_int_float_false.cu
     src/nn/specializations/knn.cu
   )
   set_target_properties(raft_nn_lib PROPERTIES OUTPUT_NAME raft_nn)

cpp/src/nn/specializations/fused_l2_knn_int_float_false.cu

@@ -0,0 +1,41 @@
+/*
+ * Copyright (c) 2022, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <cstdint>
+#include <raft/spatial/knn/detail/fused_l2_knn.cuh>
+
+namespace raft {
+namespace spatial {
+namespace knn {
+namespace detail {
+
+template void fusedL2Knn<int, float, false>(size_t D,
+                                            int* out_inds,
+                                            float* out_dists,
+                                            const float* index,
+                                            const float* query,
+                                            size_t n_index_rows,
+                                            size_t n_query_rows,
+                                            int k,
+                                            bool rowMajorIndex,
+                                            bool rowMajorQuery,
+                                            cudaStream_t stream,
+                                            raft::distance::DistanceType metric);
+
+};  // namespace detail
+};  // namespace knn
+};  // namespace spatial
+};  // namespace raft

cpp/src/nn/specializations/fused_l2_knn_int_float_true.cu

@@ -0,0 +1,40 @@
+/*
+ * Copyright (c) 2022, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <cstdint>
+#include <raft/spatial/knn/detail/fused_l2_knn.cuh>
+
+namespace raft {
+namespace spatial {
+namespace knn {
+namespace detail {
+template void fusedL2Knn<int, float, true>(size_t D,
+                                           int* out_inds,
+                                           float* out_dists,
+                                           const float* index,
+                                           const float* query,
+                                           size_t n_index_rows,
+                                           size_t n_query_rows,
+                                           int k,
+                                           bool rowMajorIndex,
+                                           bool rowMajorQuery,
+                                           cudaStream_t stream,
+                                           raft::distance::DistanceType metric);
+
+};  // namespace detail
+};  // namespace knn
+};  // namespace spatial
+};  // namespace raft

cpp/src/nn/specializations/fused_l2_knn_long_float_false.cu

@@ -0,0 +1,40 @@
+/*
+ * Copyright (c) 2022, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <cstdint>
+#include <raft/spatial/knn/detail/fused_l2_knn.cuh>
+
+namespace raft {
+namespace spatial {
+namespace knn {
+namespace detail {
+
+template void fusedL2Knn<long, float, false>(size_t D,
+                                             long* out_inds,
+                                             float* out_dists,
+                                             const float* index,
+                                             const float* query,
+                                             size_t n_index_rows,
+                                             size_t n_query_rows,
+                                             int k,
+                                             bool rowMajorIndex,
+                                             bool rowMajorQuery,
+                                             cudaStream_t stream,
+                                             raft::distance::DistanceType metric);
+};  // namespace detail
+};  // namespace knn
+};  // namespace spatial
+};  // namespace raft

cpp/src/nn/specializations/fused_l2_knn_long_float_true.cu

@@ -0,0 +1,40 @@
+/*
+ * Copyright (c) 2022, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <cstdint>
+#include <raft/spatial/knn/detail/fused_l2_knn.cuh>
+
+namespace raft {
+namespace spatial {
+namespace knn {
+namespace detail {
+
+template void fusedL2Knn<long, float, true>(size_t D,
+                                            long* out_inds,
+                                            float* out_dists,
+                                            const float* index,
+                                            const float* query,
+                                            size_t n_index_rows,
+                                            size_t n_query_rows,
+                                            int k,
+                                            bool rowMajorIndex,
+                                            bool rowMajorQuery,
+                                            cudaStream_t stream,
+                                            raft::distance::DistanceType metric);
+};  // namespace detail
+};  // namespace knn
+};  // namespace spatial
+};  // namespace raft