Pass store_dataset argument along to cuVS CAGRA (#4173)

divyegala · facebook-github-bot · commit 6b652892ff3a · 2025-02-20T19:30:30.000-08:00
Summary: This is required to enable lazy setting of a device copy of the training dataset to a cuVS CAGRA index. Pull Request resolved: #4173 Reviewed By: mnorris11 Differential Revision: D69795662 Pulled By: gtwang01 fbshipit-source-id: 68cda198ed7983800b64d3e5fac1b77ff55ecd12
diff --git a/faiss/gpu/impl/CuvsCagra.cu b/faiss/gpu/impl/CuvsCagra.cu
@@ -69,6 +69,7 @@ CuvsCagra::CuvsCagra(
 
     index_params_.intermediate_graph_degree = intermediate_graph_degree;
     index_params_.graph_degree = graph_degree;
+    index_params_.attach_dataset_on_build = store_dataset;
 
     if (!ivf_pq_search_params_) {
         ivf_pq_search_params_ =
@@ -243,6 +244,7 @@ void CuvsCagra::search(
                     storage_, n_, dim_);
             cuvs_index->update_dataset(raft_handle, dataset);
         }
+        store_dataset_ = true;
     }
 
     auto queries_view = raft::make_device_matrix_view<const float, int64_t>(

Original file line number	Diff line number	Diff line change
`@@ -69,6 +69,7 @@ CuvsCagra::CuvsCagra(`
`69`	`69`
`70`	`70`	`index_params_.intermediate_graph_degree = intermediate_graph_degree;`
`71`	`71`	`index_params_.graph_degree = graph_degree;`
	`72`	`+ index_params_.attach_dataset_on_build = store_dataset;`
`72`	`73`
`73`	`74`	`if (!ivf_pq_search_params_) {`
`74`	`75`	`ivf_pq_search_params_ =`
`@@ -243,6 +244,7 @@ void CuvsCagra::search(`
`243`	`244`	`storage_, n_, dim_);`
`244`	`245`	`cuvs_index->update_dataset(raft_handle, dataset);`
`245`	`246`	`}`
	`247`	`+ store_dataset_ = true;`
`246`	`248`	`}`
`247`	`249`
`248`	`250`	`auto queries_view = raft::make_device_matrix_view<const float, int64_t>(`