Adding limit-tests to build.sh. Removing default template args per

reviews

build.sh

@@ -19,7 +19,7 @@ ARGS=$*
 REPODIR=$(cd $(dirname $0); pwd)
 
 VALIDARGS="clean libraft pylibraft raft-dask docs tests bench clean -v -g --install --compile-libs --compile-nn --compile-dist --allgpuarch --no-nvtx --show_depr_warn -h --buildfaiss --minimal-deps"
-HELP="$0 [<target> ...] [<flag> ...] [--cmake-args=\"<args>\"] [--cache-tool=<tool>]
+HELP="$0 [<target> ...] [<flag> ...] [--cmake-args=\"<args>\"] [--cache-tool=<tool>] [--limit-tests=<targets>]
  where <target> is:
    clean            - remove all existing build artifacts and configuration (start over)
    libraft          - build the raft C++ code only. Also builds the C-wrapper library
@@ -40,6 +40,7 @@ HELP="$0 [<target> ...] [<flag> ...] [--cmake-args=\"<args>\"] [--cache-tool=<to
                                   the only option to be supported)
    --minimal-deps              - disables dependencies like thrust so they can be overridden.
                                  can be useful for a pure header-only install
+   --limit-tests               - semicolon-separated list of test executables to compile (e.g. TEST_SPATIAL;TEST_CLUSTER)
    --allgpuarch                - build for all supported GPU architectures
    --buildfaiss                - build faiss statically into raft
    --install                   - install cmake targets
@@ -50,7 +51,7 @@ HELP="$0 [<target> ...] [<flag> ...] [--cmake-args=\"<args>\"] [--cache-tool=<to
                                  to speedup the build process.
    -h                          - print this text
 
- default action (no args) is to build both libraft and raft-dask targets
+ default action (no args) is to build libraft, tests, pylibraft and raft-dask targets
 "
 LIBRAFT_BUILD_DIR=${LIBRAFT_BUILD_DIR:=${REPODIR}/cpp/build}
 SPHINX_BUILD_DIR=${REPODIR}/docs
@@ -70,6 +71,8 @@ COMPILE_NN_LIBRARY=OFF
 COMPILE_DIST_LIBRARY=OFF
 ENABLE_NN_DEPENDENCIES=OFF
 
+TEST_TARGETS="CLUSTER_TEST;CORE_TEST;DISTANCE_TEST;LABEL_TEST;LINALG_TEST;MATRIX_TEST;RANDOM_TEST;SOLVERS_TEST;SPARSE_TEST;SPARSE_DIST_TEST;SPARSE_NN_TEST;SPATIAL_TEST;STATS_TEST;UTILS_TEST"
+
 ENABLE_thrust_DEPENDENCY=ON
 
 CACHE_ARGS=""
@@ -136,6 +139,21 @@ function cacheTool {
     fi
 }
 
+function limitTests {
+    # Check for option to limit the set of test binaries to build
+    if [[ -n $(echo $ARGS | { grep -E "\-\-limit\-tests" || true; } ) ]]; then
+        # There are possible weird edge cases that may cause this regex filter to output nothing and fail silently
+        # the true pipe will catch any weird edge cases that may happen and will cause the program to fall back
+        # on the invalid option error
+        LIMIT_TEST_TARGETS=$(echo $ARGS | sed -e 's/.*--limit-tests=//' -e 's/ .*//')
+        if [[ -n ${LIMIT_TEST_TARGETS} ]]; then
+            # Remove the full LIMIT_TEST_TARGETS argument from list of args so that it passes validArgs function
+            ARGS=${ARGS//--limit-tests=$LIMIT_TEST_TARGETS/}
+            TEST_TARGETS=${LIMIT_TEST_TARGETS}
+        fi
+    fi
+}
+
 if hasArg -h || hasArg --help; then
     echo "${HELP}"
     exit 0
@@ -145,6 +163,7 @@ fi
 if (( ${NUMARGS} != 0 )); then
     cmakeArgs
     cacheTool
+    limitTests
     for a in ${ARGS}; do
         if ! (echo " ${VALIDARGS} " | grep -q " ${a} "); then
             echo "Invalid option: ${a}"
@@ -194,7 +213,7 @@ if hasArg tests || (( ${NUMARGS} == 0 )); then
     COMPILE_DIST_LIBRARY=ON
     ENABLE_NN_DEPENDENCIES=ON
     COMPILE_NN_LIBRARY=ON
-    CMAKE_TARGET="${CMAKE_TARGET};CLUSTER_TEST;CORE_TEST;DISTANCE_TEST;LABEL_TEST;LINALG_TEST;MATRIX_TEST;RANDOM_TEST;SOLVERS_TEST;SPARSE_TEST;SPARSE_DIST_TEST;SPARSE_NN_TEST;SPATIAL_TEST;STATS_TEST;UTILS_TEST"
+    CMAKE_TARGET="${CMAKE_TARGET};${TEST_TARGETS}"
 fi
 
 if hasArg bench || (( ${NUMARGS} == 0 )); then

cpp/include/raft/spatial/knn/ball_cover.cuh

@@ -39,10 +39,7 @@ namespace knn {
  * @param[in] handle library resource management handle
  * @param[inout] index an empty (and not previous built) instance of BallCoverIndex
  */
-template <typename idx_t,
-          typename value_t,
-          typename int_t,
-          typename matrix_idx_t>
+template <typename idx_t, typename value_t, typename int_t, typename matrix_idx_t>
 void rbc_build_index(const raft::handle_t& handle,
                      BallCoverIndex<idx_t, value_t, int_t, matrix_idx_t>& index)
 {
@@ -87,10 +84,7 @@ void rbc_build_index(const raft::handle_t& handle,
  *               many datasets can still have great recall even by only
  *               looking in the closest landmark.
  */
-template <typename idx_t,
-          typename value_t,
-          typename int_t,
-          typename matrix_idx_t>
+template <typename idx_t, typename value_t, typename int_t, typename matrix_idx_t>
 void rbc_all_knn_query(const raft::handle_t& handle,
                        BallCoverIndex<idx_t, value_t, int_t, matrix_idx_t>& index,
                        int_t k,

cpp/include/raft/spatial/knn/ivf_flat.cuh

@@ -100,19 +100,16 @@ inline auto build(
  *
  * @return the constructed ivf-flat index
  */
-template <typename value_t,
-          typename idx_t,
-          typename int_t,
-          typename matrix_idx_t>
+template <typename value_t, typename idx_t>
 auto build_index(const handle_t& handle,
-                 raft::device_matrix_view<const value_t, matrix_idx_t, row_major> dataset,
+                 raft::device_matrix_view<const value_t, idx_t, row_major> dataset,
                  const index_params& params) -> index<value_t, idx_t>
 {
   return raft::spatial::knn::ivf_flat::detail::build(handle,
                                                      params,
                                                      dataset.data_handle(),
                                                      static_cast<idx_t>(dataset.extent(0)),
-                                                     static_cast<int_t>(dataset.extent(1)));
+                                                     static_cast<idx_t>(dataset.extent(1)));
 }
 
 /**
@@ -191,15 +188,12 @@ inline auto extend(const handle_t& handle,
  *
  * @return the constructed extended ivf-flat index
  */
-template <typename value_t,
-          typename idx_t,
-          typename int_t        = std::uint32_t,
-          typename matrix_idx_t = std::uint32_t>
+template <typename value_t, typename idx_t>
 auto extend(const handle_t& handle,
             const index<value_t, idx_t>& orig_index,
-            raft::device_matrix_view<const value_t, matrix_idx_t, row_major> new_vectors,
-            std::optional<raft::device_vector_view<const idx_t, matrix_idx_t>> new_indices =
-              std::nullopt) -> index<value_t, idx_t>
+            raft::device_matrix_view<const value_t, idx_t, row_major> new_vectors,
+            std::optional<raft::device_vector_view<const idx_t, idx_t>> new_indices = std::nullopt)
+  -> index<value_t, idx_t>
 {
   return raft::spatial::knn::ivf_flat::detail::extend<value_t, idx_t>(
     handle,
@@ -248,15 +242,11 @@ inline void extend(const handle_t& handle,
  *    If the original index is empty (`orig_index.size() == 0`), you can pass `std::nullopt`
  *    here to imply a continuous range `[0...n_rows)`.
  */
-template <typename value_t,
-          typename idx_t,
-          typename int_t        = std::uint32_t,
-          typename matrix_idx_t = std::uint32_t>
-void extend(
-  const handle_t& handle,
-  index<value_t, idx_t>* index,
-  raft::device_matrix_view<const value_t, matrix_idx_t, row_major> new_vectors,
-  std::optional<raft::device_vector_view<const idx_t, matrix_idx_t>> new_indices = std::nullopt)
+template <typename value_t, typename idx_t>
+void extend(const handle_t& handle,
+            index<value_t, idx_t>* index,
+            raft::device_matrix_view<const value_t, idx_t, row_major> new_vectors,
+            std::optional<raft::device_vector_view<const idx_t, idx_t>> new_indices = std::nullopt)
 {
   *index = extend(handle,
                   *index,
@@ -363,25 +353,22 @@ inline void search(const handle_t& handle,
  * @param[in] params configure the search
  * @param[in] k the number of neighbors to find for each query.
  */
-template <typename value_t,
-          typename idx_t,
-          typename int_t        = std::uint32_t,
-          typename matrix_idx_t = std::uint32_t>
+template <typename value_t, typename idx_t, typename int_t>
 void search(const handle_t& handle,
             const index<value_t, idx_t>& index,
-            raft::device_matrix_view<const value_t, matrix_idx_t, row_major> queries,
-            raft::device_matrix_view<idx_t, matrix_idx_t, row_major> neighbors,
-            raft::device_matrix_view<idx_t, matrix_idx_t, float> distances,
+            raft::device_matrix_view<const value_t, idx_t, row_major> queries,
+            raft::device_matrix_view<idx_t, idx_t, row_major> neighbors,
+            raft::device_matrix_view<idx_t, idx_t, float> distances,
             const search_params& params,
             int_t k)
 {
   RAFT_EXPECTS(
     queries.extent(0) == neighbors.extent(0) && queries.extent(0) == distances.extent(0),
     "Number of rows in output neighbors and distances matrices must equal the number of queries.");
 
-  RAFT_EXPECTS(neighbors.extent(1) == distances.extent(1) &&
-                 neighbors.extent(1) == static_cast<matrix_idx_t>(k),
-               "Number of columns in output neighbors and distances matrices must equal k");
+  RAFT_EXPECTS(
+    neighbors.extent(1) == distances.extent(1) && neighbors.extent(1) == static_cast<idx_t>(k),
+    "Number of columns in output neighbors and distances matrices must equal k");
 
   RAFT_EXPECTS(queries.extent(1) == index.dim(),
                "Number of query dimensions should equal number of dimensions in the index.");

cpp/test/spatial/ann_ivf_flat.cu

@@ -208,10 +208,10 @@ class AnnIVFFlatTest : public ::testing::TestWithParam<AnnIvfFlatInputs> {
         index_params.add_data_on_build        = false;
         index_params.kmeans_trainset_fraction = 0.5;
 
-        auto database_view = raft::make_device_matrix_view<const DataT>(
+        auto database_view = raft::make_device_matrix_view<const DataT, IdxT>(
           (const DataT*)database.data(), ps.num_db_vecs, ps.dim);
 
-        auto index = ivf_flat::build_index<DataT, IdxT>(handle_, database_view, index_params);
+        auto index = ivf_flat::build_index(handle_, database_view, index_params);
 
         rmm::device_uvector<IdxT> vector_indices(ps.num_db_vecs, stream_);
         thrust::sequence(handle_.get_thrust_policy(),
@@ -221,16 +221,16 @@ class AnnIVFFlatTest : public ::testing::TestWithParam<AnnIvfFlatInputs> {
 
         int64_t half_of_data = ps.num_db_vecs / 2;
 
-        auto half_of_data_view = raft::make_device_matrix_view<const DataT, int>(
+        auto half_of_data_view = raft::make_device_matrix_view<const DataT, IdxT>(
           (const DataT*)database.data(), half_of_data, ps.dim);
 
-        auto index_2 = ivf_flat::extend<DataT, IdxT>(handle_, index, half_of_data_view);
+        auto index_2 = ivf_flat::extend(handle_, index, half_of_data_view);
 
-        ivf_flat::extend<DataT, IdxT>(handle_,
-                                      &index_2,
-                                      database.data() + half_of_data * ps.dim,
-                                      vector_indices.data() + half_of_data,
-                                      int64_t(ps.num_db_vecs) - half_of_data);
+        ivf_flat::extend(handle_,
+                         &index_2,
+                         database.data() + half_of_data * ps.dim,
+                         vector_indices.data() + half_of_data,
+                         int64_t(ps.num_db_vecs) - half_of_data);
 
         ivf_flat::search(handle_,
                          search_params,