Add tests for iterative_build_graph_index

rapidsai · Jan 28, 2025 · ce17775 · ce17775
1 parent 217f226
commit ce17775
Showing 1 changed file with 23 additions and 5 deletions.
diff --git a/cpp/tests/neighbors/ann_cagra.cuh b/cpp/tests/neighbors/ann_cagra.cuh
@@ -254,6 +254,8 @@ enum class graph_build_algo {
   IVF_PQ,
   /* Experimental, use NN-Descent to build all-neighbors knn graph */
   NN_DESCENT,
+  /* Experimental, iteratively execute CAGRA's search() and optimize() */
+  ITERATIVE_CAGRA_SEARCH,
   /* Choose default automatically */
   AUTO
 };
@@ -294,7 +296,7 @@ inline ::std::ostream& operator<<(::std::ostream& os, const AnnCagraInputs& p)
   };
 
   std::vector<std::string> algo       = {"single-cta", "multi_cta", "multi_kernel", "auto"};
-  std::vector<std::string> build_algo = {"IVF_PQ", "NN_DESCENT", "AUTO"};
+  std::vector<std::string> build_algo = {"IVF_PQ", "NN_DESCENT", "ITERATIVE_CAGRA_SEARCH", "AUTO"};
   os << "{n_queries=" << p.n_queries << ", dataset shape=" << p.n_rows << "x" << p.dim
      << ", k=" << p.k << ", " << algo.at((int)p.algo) << ", max_queries=" << p.max_queries
      << ", itopk_size=" << p.itopk_size << ", search_width=" << p.search_width
@@ -374,6 +376,10 @@ class AnnCagraTest : public ::testing::TestWithParam<AnnCagraInputs> {
               index_params.intermediate_graph_degree, index_params.metric);
             break;
           }
+          case graph_build_algo::ITERATIVE_CAGRA_SEARCH: {
+            index_params.graph_build_params = graph_build_params::iterative_search_params();
+            break;
+          }
           case graph_build_algo::AUTO:
             // do nothing
             break;
@@ -551,6 +557,10 @@ class AnnCagraAddNodesTest : public ::testing::TestWithParam<AnnCagraInputs> {
               graph_build_params::nn_descent_params(index_params.intermediate_graph_degree);
             break;
           }
+          case graph_build_algo::ITERATIVE_CAGRA_SEARCH: {
+            index_params.graph_build_params = graph_build_params::iterative_search_params();
+            break;
+          }
           case graph_build_algo::AUTO:
             // do nothing
             break;
@@ -752,6 +762,10 @@ class AnnCagraFilterTest : public ::testing::TestWithParam<AnnCagraInputs> {
               graph_build_params::nn_descent_params(index_params.intermediate_graph_degree);
             break;
           }
+          case graph_build_algo::ITERATIVE_CAGRA_SEARCH: {
+            index_params.graph_build_params = graph_build_params::iterative_search_params();
+            break;
+          }
           case graph_build_algo::AUTO:
             // do nothing
             break;
@@ -882,7 +896,7 @@ inline std::vector<AnnCagraInputs> generate_inputs()
     {1000},
     {1, 8, 17},
     {16},  // k
-    {graph_build_algo::NN_DESCENT},
+    {graph_build_algo::NN_DESCENT, graph_build_algo::ITERATIVE_CAGRA_SEARCH},
     {search_algo::SINGLE_CTA, search_algo::MULTI_CTA, search_algo::MULTI_KERNEL},
     {0, 10},  // query size
     {0},
@@ -917,7 +931,7 @@ inline std::vector<AnnCagraInputs> generate_inputs()
     {1000},
     {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},
+    {graph_build_algo::IVF_PQ, graph_build_algo::NN_DESCENT, graph_build_algo::ITERATIVE_CAGRA_SEARCH},
     {search_algo::AUTO},
     {10},
     {0},
@@ -935,7 +949,9 @@ inline std::vector<AnnCagraInputs> generate_inputs()
     {1000},
     {64},
     {16},
-    {graph_build_algo::IVF_PQ, graph_build_algo::NN_DESCENT},
+    {graph_build_algo::IVF_PQ,
+     graph_build_algo::NN_DESCENT,
+     graph_build_algo::ITERATIVE_CAGRA_SEARCH},
     {search_algo::AUTO},
     {10},
     {0, 8, 16, 32},  // team_size
@@ -953,7 +969,9 @@ inline std::vector<AnnCagraInputs> generate_inputs()
     {1000},
     {64},
     {16},
-    {graph_build_algo::IVF_PQ, graph_build_algo::NN_DESCENT},
+    {graph_build_algo::IVF_PQ,
+     graph_build_algo::NN_DESCENT,
+     graph_build_algo::ITERATIVE_CAGRA_SEARCH},
     {search_algo::AUTO},
     {10},
     {0},  // team_size