Splitting fused l2 knn specializations

README.md

@@ -62,7 +62,7 @@ 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.
 

build.sh

@@ -18,7 +18,7 @@ ARGS=$*
 # script, and that this script resides in the repo dir!
 REPODIR=$(cd $(dirname $0); pwd)
 
-VALIDARGS="clean libraft pyraft docs -v -g --compilelibs --allgpuarch --nvtx --show_depr_warn -h --buildgtest --buildfaiss"
+VALIDARGS="clean libraft pyraft docs -v -g --compilelibs --allgpuarch --nvtx --show_depr_warn -h --nogtest --buildfaiss"
 HELP="$0 [<target> ...] [<flag> ...]
  where <target> is:
    clean            - remove all existing build artifacts and configuration (start over)
@@ -33,6 +33,7 @@ HELP="$0 [<target> ...] [<flag> ...]
    --compilelibs    - compile shared libraries
    --allgpuarch     - build for all supported GPU architectures
    --buildfaiss     - build faiss statically into raft
+   --nogtest        - do not build google tests for libraft
    --nvtx           - Enable nvtx for profiling support
    --show_depr_warn - show cmake deprecation warnings
    -h               - print this text
@@ -49,14 +50,15 @@ BUILD_DIRS="${CPP_RAFT_BUILD_DIR} ${PY_RAFT_BUILD_DIR} ${PYTHON_DEPS_CLONE}"
 CMAKE_LOG_LEVEL=""
 VERBOSE_FLAG=""
 BUILD_ALL_GPU_ARCH=0
-BUILD_TESTS=ON
+BUILD_TESTS=YES
 BUILD_STATIC_FAISS=OFF
-COMPILE_LIBRARIES=OFF
-ENABLE_NN_DEPENDENCIES=ON
+COMPILE_LIBRARIES=${BUILD_TESTS}
+ENABLE_NN_DEPENDENCIES=${BUILD_TESTS}
 SINGLEGPU=""
 NVTX=OFF
 CLEAN=0
-BUILD_DISABLE_DEPRECATION_WARNING=ON
+DISABLE_DEPRECATION_WARNINGS=ON
+CMAKE_TARGET=""
 
 # Set defaults for vars that may not have been defined externally
 #  FIXME: if INSTALL_PREFIX is not set, check PREFIX, then check
@@ -97,18 +99,20 @@ if hasArg -g; then
     BUILD_TYPE=Debug
 fi
 
-if hasArg --compilelibs; then
-    COMPILE_LIBRARIES=ON
-fi
-
 if hasArg --allgpuarch; then
     BUILD_ALL_GPU_ARCH=1
 fi
-if hasArg --buildgtest; then
-    BUILD_GTEST=ON
+if hasArg --nogtest; then
+    BUILD_TESTS=OFF
+    COMPILE_LIBRARIES=OFF
+    ENABLE_NN_DEPENDENCIES=OFF
+fi
+if hasArg --compilelibs; then
+    COMPILE_LIBRARIES=ON
+    ENABLE_NN_DEPENDENCIES=ON
 fi
 if hasArg --buildfaiss; then
-      BUILD_STATIC_FAISS=ON
+    BUILD_STATIC_FAISS=ON
 fi
 if hasArg --singlegpu; then
     SINGLEGPU="--singlegpu"
@@ -117,7 +121,7 @@ if hasArg --nvtx; then
     NVTX=ON
 fi
 if hasArg --show_depr_warn; then
-    BUILD_DISABLE_DEPRECATION_WARNING=OFF
+    DISABLE_DEPRECATION_WARNINGS=OFF
 fi
 if hasArg clean; then
     CLEAN=1
@@ -160,14 +164,14 @@ if (( ${NUMARGS} == 0 )) || hasArg libraft || hasArg docs; then
           -DRAFT_COMPILE_LIBRARIES=${COMPILE_LIBRARIES} \
           -DRAFT_ENABLE_NN_DEPENDENCIES=${ENABLE_NN_DEPENDENCIES} \
           -DNVTX=${NVTX} \
-          -DDISABLE_DEPRECATION_WARNING=${BUILD_DISABLE_DEPRECATION_WARNING} \
+          -DDISABLE_DEPRECATION_WARNINGS=${DISABLE_DEPRECATION_WARNINGS} \
           -DBUILD_TESTS=${BUILD_TESTS} \
           -DRAFT_USE_FAISS_STATIC=${BUILD_STATIC_FAISS} \
           ..
 
   if (( ${NUMARGS} == 0 )) || hasArg libraft; then
       # Run all c++ targets at once
-      cmake --build  ${CPP_RAFT_BUILD_DIR} -j${PARALLEL_LEVEL} ${VERBOSE_FLAG}
+      cmake --build  ${CPP_RAFT_BUILD_DIR} -j${PARALLEL_LEVEL} ${VERBOSE_FLAG} ${CMAKE_TARGET}
   fi
 fi
 

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