Reduce CAGRA test runtime (#602)

Currently, running `NEIGHBORS_ANN_CAGRA_TEST` takes: [0.96 hours on CUDA 11.8, V100 (x86)](https://github.com/rapidsai/cuvs/actions/runs/12913409417/job/36012418022?pr=596#step:8:1718) [1.59 hours on CUDA 12.5, V100 (x86)](https://github.com/rapidsai/cuvs/actions/runs/12913409417/job/36012418329?pr=596#step:8:492) [0.28 hours on CUDA 12.0, A100 (ARM)](https://github.com/rapidsai/cuvs/actions/runs/12913409417/job/36012418741?pr=596#step:8:1729) Individual tests should be able to complete in less than an hour. Ideally, less than 10 minutes. This PR proposes some changes to CAGRA tests: - Each CAGRA type is now its own test executable (e.g. `NEIGHBORS_ANN_CAGRA_FLOAT_UINT32_TEST`) - Some parameter combinations were trimmed by ~50% Authors: - Bradley Dice (https://github.com/bdice) - Tamas Bela Feher (https://github.com/tfeher) - Artem M. Chirkin (https://github.com/achirkin) Approvers: - Artem M. Chirkin (https://github.com/achirkin) - Divye Gala (https://github.com/divyegala) URL: #602
rapidsai · Jan 27, 2025 · 5527cdf · 5527cdf
1 parent 11ecbb2
commit 5527cdf
Show file tree

Hide file tree

Showing 5 changed files with 208 additions and 29 deletions.
diff --git a/cpp/test/CMakeLists.txt b/cpp/test/CMakeLists.txt
@@ -138,19 +138,28 @@ if(BUILD_TESTS)
   )
 
   ConfigureTest(
-    NAME
-    NEIGHBORS_ANN_CAGRA_TEST
-    PATH
-    neighbors/ann_cagra/bug_extreme_inputs_oob.cu
-    neighbors/ann_cagra/bug_multi_cta_crash.cu
-    neighbors/ann_cagra/test_float_uint32_t.cu
-    neighbors/ann_cagra/test_half_uint32_t.cu
-    neighbors/ann_cagra/test_int8_t_uint32_t.cu
-    neighbors/ann_cagra/test_uint8_t_uint32_t.cu
-    GPUS
-    1
-    PERCENT
-    100
+    NAME NEIGHBORS_ANN_CAGRA_TEST_BUGS PATH neighbors/ann_cagra/bug_extreme_inputs_oob.cu
+    neighbors/ann_cagra/bug_multi_cta_crash.cu GPUS 1 PERCENT 100
+  )
+
+  ConfigureTest(
+    NAME NEIGHBORS_ANN_CAGRA_FLOAT_UINT32_TEST PATH neighbors/ann_cagra/test_float_uint32_t.cu GPUS
+    1 PERCENT 100
+  )
+
+  ConfigureTest(
+    NAME NEIGHBORS_ANN_CAGRA_HALF_UINT32_TEST PATH neighbors/ann_cagra/test_half_uint32_t.cu GPUS 1
+    PERCENT 100
+  )
+
+  ConfigureTest(
+    NAME NEIGHBORS_ANN_CAGRA_INT8_UINT32_TEST PATH neighbors/ann_cagra/test_int8_t_uint32_t.cu GPUS
+    1 PERCENT 100
+  )
+
+  ConfigureTest(
+    NAME NEIGHBORS_ANN_CAGRA_UINT8_UINT32_TEST PATH neighbors/ann_cagra/test_uint8_t_uint32_t.cu
+    GPUS 1 PERCENT 100
   )
 
   ConfigureTest(

diff --git a/cpp/test/neighbors/ann_cagra.cuh b/cpp/test/neighbors/ann_cagra.cuh
@@ -876,14 +876,15 @@ class AnnCagraFilterTest : public ::testing::TestWithParam<AnnCagraInputs> {
 inline std::vector<AnnCagraInputs> generate_inputs()
 {
   // TODO(tfeher): test MULTI_CTA kernel with search_width > 1 to allow multiple CTA per queries
+  // Charge graph dim, search algo and max_query parameter
   std::vector<AnnCagraInputs> inputs = raft::util::itertools::product<AnnCagraInputs>(
     {100},
     {1000},
     {1, 8, 17},
-    {1, 16},  // k
-    {graph_build_algo::IVF_PQ, graph_build_algo::NN_DESCENT},
+    {16},  // k
+    {graph_build_algo::NN_DESCENT},
     {search_algo::SINGLE_CTA, search_algo::MULTI_CTA, search_algo::MULTI_KERNEL},
-    {0, 1, 10, 100},  // query size
+    {0, 10},  // query size
     {0},
     {256},
     {1},
@@ -892,11 +893,30 @@ inline std::vector<AnnCagraInputs> generate_inputs()
     {true},
     {0.995});
 
+  // Fixed dim, and changing neighbors and query size (output matrix size)
   auto inputs2 = raft::util::itertools::product<AnnCagraInputs>(
+    {1, 100},
+    {1000},
+    {8},
+    {1, 16},  // k
+    {graph_build_algo::NN_DESCENT},
+    {search_algo::SINGLE_CTA, search_algo::MULTI_CTA, search_algo::MULTI_KERNEL},
+    {0},  // query size
+    {0},
+    {256},
+    {1},
+    {cuvs::distance::DistanceType::L2Expanded, cuvs::distance::DistanceType::InnerProduct},
+    {false},
+    {true},
+    {0.995});
+  inputs.insert(inputs.end(), inputs2.begin(), inputs2.end());
+
+  // Varying dim and build algo.
+  inputs2 = raft::util::itertools::product<AnnCagraInputs>(
     {100},
     {1000},
-    {1, 3, 5, 7, 8, 17, 64, 128, 137, 192, 256, 512, 619, 1024},  // dim
-    {16},                                                         // k
+    {1, 3, 5, 7, 8, 17, 64, 128, 137, 192, 256, 512, 1024},  // dim
+    {16},                                                    // k
     {graph_build_algo::IVF_PQ, graph_build_algo::NN_DESCENT},
     {search_algo::AUTO},
     {10},
@@ -908,6 +928,8 @@ inline std::vector<AnnCagraInputs> generate_inputs()
     {true},
     {0.995});
   inputs.insert(inputs.end(), inputs2.begin(), inputs2.end());
+
+  // Varying team_size, graph_build_algo
   inputs2 = raft::util::itertools::product<AnnCagraInputs>(
     {100},
     {1000},
@@ -925,6 +947,7 @@ inline std::vector<AnnCagraInputs> generate_inputs()
     {0.995});
   inputs.insert(inputs.end(), inputs2.begin(), inputs2.end());
 
+  // Varying graph_build_algo, itopk_size
   inputs2 = raft::util::itertools::product<AnnCagraInputs>(
     {100},
     {1000},
@@ -942,9 +965,10 @@ inline std::vector<AnnCagraInputs> generate_inputs()
     {0.995});
   inputs.insert(inputs.end(), inputs2.begin(), inputs2.end());
 
+  // Varying n_rows, host_dataset
   inputs2 = raft::util::itertools::product<AnnCagraInputs>(
     {100},
-    {10000, 20000},
+    {10000},
     {32},
     {10},
     {graph_build_algo::AUTO},
@@ -959,7 +983,8 @@ inline std::vector<AnnCagraInputs> generate_inputs()
     {0.985});
   inputs.insert(inputs.end(), inputs2.begin(), inputs2.end());
 
-  // a few PQ configurations
+  // A few PQ configurations.
+  // Varying dim, vq_n_centers
   inputs2 = raft::util::itertools::product<AnnCagraInputs>(
     {100},
     {10000},
@@ -987,7 +1012,8 @@ inline std::vector<AnnCagraInputs> generate_inputs()
     }
   }
 
-  // refinement options
+  // Refinement options
+  // Varying host_dataset, ivf_pq_search_refine_ratio
   inputs2 = raft::util::itertools::product<AnnCagraInputs>(
     {100},
     {5000},
@@ -1006,10 +1032,11 @@ inline std::vector<AnnCagraInputs> generate_inputs()
     {1.0f, 2.0f, 3.0f});
   inputs.insert(inputs.end(), inputs2.begin(), inputs2.end());
 
+  // Varying dim, adding non_owning_memory_buffer_flag
   inputs2 = raft::util::itertools::product<AnnCagraInputs>(
     {100},
     {1000},
-    {1, 3, 5, 7, 8, 17, 64, 128, 137, 192, 256, 512, 619, 1024},  // dim
+    {1, 5, 8, 64, 137, 256, 619, 1024},  // dim
     {10},
     {graph_build_algo::IVF_PQ},
     {search_algo::AUTO},
@@ -1029,6 +1056,143 @@ inline std::vector<AnnCagraInputs> generate_inputs()
   return inputs;
 }
 
-const std::vector<AnnCagraInputs> inputs = generate_inputs();
+inline std::vector<AnnCagraInputs> generate_addnode_inputs()
+{
+  // changing dim
+  std::vector<AnnCagraInputs> inputs = raft::util::itertools::product<AnnCagraInputs>(
+    {100},
+    {1000},
+    {1, 8, 17, 64, 128, 137, 512, 1024},  // dim
+    {16},                                 // k
+    {graph_build_algo::NN_DESCENT},
+    {search_algo::AUTO},
+    {10},
+    {0},
+    {64},
+    {1},
+    {cuvs::distance::DistanceType::L2Expanded, cuvs::distance::DistanceType::InnerProduct},
+    {false},
+    {true},
+    {0.995});
+
+  // testing host and device datasets
+  auto inputs2 = raft::util::itertools::product<AnnCagraInputs>(
+    {100},
+    {10000},
+    {32},
+    {10},
+    {graph_build_algo::AUTO},
+    {search_algo::AUTO},
+    {10},
+    {0},  // team_size
+    {64},
+    {1},
+    {cuvs::distance::DistanceType::L2Expanded, cuvs::distance::DistanceType::InnerProduct},
+    {false, true},
+    {false},
+    {0.985});
+  inputs.insert(inputs.end(), inputs2.begin(), inputs2.end());
+
+  // a few PQ configurations
+  inputs2 = raft::util::itertools::product<AnnCagraInputs>(
+    {100},
+    {10000},
+    {192, 1024},  // dim
+    {16},         // k
+    {graph_build_algo::IVF_PQ},
+    {search_algo::AUTO},
+    {10},
+    {0},
+    {64},
+    {1},
+    {cuvs::distance::DistanceType::L2Expanded},
+    {false},
+    {true},
+    {0.6});                      // don't demand high recall without refinement
+  for (uint32_t pq_len : {2}) {  // for now, only pq_len = 2 is supported, more options coming soon
+    for (uint32_t vq_n_centers : {100}) {
+      for (auto input : inputs2) {
+        vpq_params ps{};
+        ps.pq_dim       = input.dim / pq_len;
+        ps.vq_n_centers = vq_n_centers;
+        input.compression.emplace(ps);
+        inputs.push_back(input);
+      }
+    }
+  }
+
+  return inputs;
+}
+
+inline std::vector<AnnCagraInputs> generate_filtering_inputs()
+{
+  // Charge graph dim, search algo
+  std::vector<AnnCagraInputs> inputs = raft::util::itertools::product<AnnCagraInputs>(
+    {100},
+    {1000},
+    {1, 8, 17},
+    {16},  // k
+    {graph_build_algo::NN_DESCENT},
+    {search_algo::SINGLE_CTA, search_algo::MULTI_CTA, search_algo::MULTI_KERNEL},
+    {0},  // query size
+    {0},
+    {256},
+    {1},
+    {cuvs::distance::DistanceType::L2Expanded, cuvs::distance::DistanceType::InnerProduct},
+    {false},
+    {true},
+    {0.995});
+
+  // Fixed dim, and changing neighbors and query size (output matrix size)
+  auto inputs2 = raft::util::itertools::product<AnnCagraInputs>(
+    {1, 100},
+    {1000},
+    {8},
+    {1, 16},  // k
+    {graph_build_algo::NN_DESCENT},
+    {search_algo::SINGLE_CTA, search_algo::MULTI_CTA, search_algo::MULTI_KERNEL},
+    {0},  // query size
+    {0},
+    {256},
+    {1},
+    {cuvs::distance::DistanceType::L2Expanded, cuvs::distance::DistanceType::InnerProduct},
+    {false},
+    {true},
+    {0.995});
+  inputs.insert(inputs.end(), inputs2.begin(), inputs2.end());
+
+  // a few PQ configurations
+  inputs2 = raft::util::itertools::product<AnnCagraInputs>(
+    {100},
+    {10000},
+    {256},  // dim
+    {16},   // k
+    {graph_build_algo::IVF_PQ},
+    {search_algo::AUTO},
+    {10},
+    {0},
+    {64},
+    {1},
+    {cuvs::distance::DistanceType::L2Expanded},
+    {false},
+    {true},
+    {0.6});                      // don't demand high recall without refinement
+  for (uint32_t pq_len : {2}) {  // for now, only pq_len = 2 is supported, more options coming soon
+    for (uint32_t vq_n_centers : {100}) {
+      for (auto input : inputs2) {
+        vpq_params ps{};
+        ps.pq_dim       = input.dim / pq_len;
+        ps.vq_n_centers = vq_n_centers;
+        input.compression.emplace(ps);
+        inputs.push_back(input);
+      }
+    }
+  }
+
+  return inputs;
+}
+const std::vector<AnnCagraInputs> inputs           = generate_inputs();
+const std::vector<AnnCagraInputs> inputs_addnode   = generate_addnode_inputs();
+const std::vector<AnnCagraInputs> inputs_filtering = generate_filtering_inputs();
 
 }  // namespace cuvs::neighbors::cagra
diff --git a/cpp/test/neighbors/ann_cagra/test_float_uint32_t.cu b/cpp/test/neighbors/ann_cagra/test_float_uint32_t.cu
@@ -32,7 +32,9 @@ TEST_P(AnnCagraFilterTestF_U32, AnnCagra) { this->testCagra(); }
 INSTANTIATE_TEST_CASE_P(AnnCagraTest, AnnCagraTestF_U32, ::testing::ValuesIn(inputs));
 INSTANTIATE_TEST_CASE_P(AnnCagraAddNodesTest,
                         AnnCagraAddNodesTestF_U32,
-                        ::testing::ValuesIn(inputs));
-INSTANTIATE_TEST_CASE_P(AnnCagraFilterTest, AnnCagraFilterTestF_U32, ::testing::ValuesIn(inputs));
+                        ::testing::ValuesIn(inputs_addnode));
+INSTANTIATE_TEST_CASE_P(AnnCagraFilterTest,
+                        AnnCagraFilterTestF_U32,
+                        ::testing::ValuesIn(inputs_filtering));
 
 }  // namespace cuvs::neighbors::cagra
diff --git a/cpp/test/neighbors/ann_cagra/test_int8_t_uint32_t.cu b/cpp/test/neighbors/ann_cagra/test_int8_t_uint32_t.cu
@@ -30,7 +30,9 @@ TEST_P(AnnCagraFilterTestI8_U32, AnnCagra) { this->testCagra(); }
 INSTANTIATE_TEST_CASE_P(AnnCagraTest, AnnCagraTestI8_U32, ::testing::ValuesIn(inputs));
 INSTANTIATE_TEST_CASE_P(AnnCagraAddNodesTest,
                         AnnCagraAddNodesTestI8_U32,
-                        ::testing::ValuesIn(inputs));
-INSTANTIATE_TEST_CASE_P(AnnCagraFilterTest, AnnCagraFilterTestI8_U32, ::testing::ValuesIn(inputs));
+                        ::testing::ValuesIn(inputs_addnode));
+INSTANTIATE_TEST_CASE_P(AnnCagraFilterTest,
+                        AnnCagraFilterTestI8_U32,
+                        ::testing::ValuesIn(inputs_filtering));
 
 }  // namespace cuvs::neighbors::cagra
diff --git a/cpp/test/neighbors/ann_cagra/test_uint8_t_uint32_t.cu b/cpp/test/neighbors/ann_cagra/test_uint8_t_uint32_t.cu
@@ -30,7 +30,9 @@ TEST_P(AnnCagraFilterTestU8_U32, AnnCagra) { this->testCagra(); }
 INSTANTIATE_TEST_CASE_P(AnnCagraTest, AnnCagraTestU8_U32, ::testing::ValuesIn(inputs));
 INSTANTIATE_TEST_CASE_P(AnnCagraAddNodesTest,
                         AnnCagraAddNodesTestU8_U32,
-                        ::testing::ValuesIn(inputs));
-INSTANTIATE_TEST_CASE_P(AnnCagraFilterTest, AnnCagraFilterTestU8_U32, ::testing::ValuesIn(inputs));
+                        ::testing::ValuesIn(inputs_addnode));
+INSTANTIATE_TEST_CASE_P(AnnCagraFilterTest,
+                        AnnCagraFilterTestU8_U32,
+                        ::testing::ValuesIn(inputs_filtering));
 
 }  // namespace cuvs::neighbors::cagra