Fix pq_compressed file created by SSD Vamana build

cpp/src/neighbors/detail/vamana/vamana_build.cuh

@@ -388,6 +388,7 @@ auto quantize_all_vectors(raft::resources const& res,
 {
   auto dim         = residuals.extent(1);
   auto vq_codebook = raft::make_device_matrix<float, uint32_t, raft::row_major>(res, 1, dim);
+  raft::linalg::map_offset(res, vq_codebook.view(), [] __device__(size_t i) { return 0; });
 
   auto codes = cuvs::neighbors::detail::process_and_fill_codes_subspaces<float, int64_t>(
     res, ps, residuals, raft::make_const_mdspan(vq_codebook.view()), pq_codebook);

examples/cpp/src/vamana_example.cu

@@ -45,7 +45,9 @@ void vamana_build_and_write(raft::device_resources const& dev_resources,
   index_params.visited_size = visited_size;
   index_params.graph_degree = degree;
   index_params.vamana_iters = iters;
-  index_params.codebooks    = vamana::deserialize_codebooks(codebook_prefix, dataset.extent(1));
+  if (codebook_prefix != "") {
+    index_params.codebooks = vamana::deserialize_codebooks(codebook_prefix, dataset.extent(1));
+  }
 
   std::cout << "Building Vamana index (search graph)" << std::endl;
 
@@ -60,20 +62,23 @@ void vamana_build_and_write(raft::device_resources const& dev_resources,
 
   std::cout << "Time to build index: " << elapsed_seconds.count() << "s\n";
 
-  // Output index to file (in-memory format)
-  serialize(dev_resources, out_fname, index);
-
-  // Output index to file (disk sector-aligned format)
-  serialize(dev_resources, out_fname + ".sector_aligned", index, false, true);
+  if (codebook_prefix != "") {
+    // Output index to file (disk sector-aligned format)
+    serialize(dev_resources, out_fname, index, false, true);
+  } else {
+    // Output index to file (in-memory format)
+    serialize(dev_resources, out_fname, index);
+  }
 }
 
 void usage()
 {
   printf(
-    "Usage: ./vamana_example <data filename> <output filename> <graph "
+    "Usage: ./vamana_example <data filename> <output filename> <datatype> <graph "
     "degree> <visited_size> <max_fraction> <iterations> <(optional) "
     "codebook prefix>\n");
-  printf("Input file expected to be binary file of fp32 vectors.\n");
+  printf("Input file expected to be binary file of vectors.\n");
+  printf("Datatype options int8 or float for this example.\n");
   printf("Graph degree sizes supported: 32, 64, 128, 256\n");
   printf("Visited_size must be > degree and a power of 2.\n");
   printf("max_fraction > 0 and <= 1. Typical values are 0.06 or 0.1.\n");
@@ -102,28 +107,44 @@ int main(int argc, char* argv[])
   // limit. raft::resource::set_workspace_to_pool_resource(dev_resources, 2 *
   // 1024 * 1024 * 1024ull);
 
-  if (argc != 7 && argc != 8) usage();
+  if (argc != 8 && argc != 9) usage();
 
   std::string data_fname      = (std::string)(argv[1]);  // Input filename
   std::string out_fname       = (std::string)argv[2];    // Output index filename
-  int degree                  = atoi(argv[3]);
-  int max_visited             = atoi(argv[4]);
-  float max_fraction          = atof(argv[5]);
-  int iters                   = atoi(argv[6]);
+  std::string dtype           = (std::string)argv[3];
+  int degree                  = atoi(argv[4]);
+  int max_visited             = atoi(argv[5]);
+  float max_fraction          = atof(argv[6]);
+  int iters                   = atoi(argv[7]);
   std::string codebook_prefix = "";
-  if (argc >= 8)
-    codebook_prefix = (std::string)argv[7];  // Path prefix to pq pivots and rotation matrix files
-
-  // Read in binary dataset file
-  auto dataset = read_bin_dataset<uint8_t, int64_t>(dev_resources, data_fname, INT_MAX);
-
-  // Simple build example to create graph and write to a file
-  vamana_build_and_write<uint8_t>(dev_resources,
+  if (argc >= 9)
+    codebook_prefix = (std::string)argv[8];  // Path prefix to pq pivots and rotation matrix files
+
+  if (dtype == "int8") {
+    // Read in binary dataset file
+    auto dataset = read_bin_dataset<int8_t, int64_t>(dev_resources, data_fname, INT_MAX);
+
+    // Simple build example to create graph and write to a file
+    vamana_build_and_write<int8_t>(dev_resources,
+                                   raft::make_const_mdspan(dataset.view()),
+                                   out_fname,
+                                   degree,
+                                   max_visited,
+                                   max_fraction,
+                                   iters,
+                                   codebook_prefix);
+  } else if (dtype == "float") {
+    // Read in binary dataset file
+    auto dataset = read_bin_dataset<float, int64_t>(dev_resources, data_fname, INT_MAX);
+
+    // Simple build example to create graph and write to a file
+    vamana_build_and_write<float>(dev_resources,
                                   raft::make_const_mdspan(dataset.view()),
                                   out_fname,
                                   degree,
                                   max_visited,
                                   max_fraction,
                                   iters,
                                   codebook_prefix);
+  }
 }