Merge branch 'main' into gpipnn_core

Author: bkarsin

.github/workflows/build.yaml

@@ -35,6 +35,20 @@ concurrency:
 permissions: {}
 
 jobs:
+  checks:
+    if: ${{ github.ref == 'refs/heads/main' }}
+    permissions:
+      actions: read
+      contents: read
+    uses: rapidsai/shared-workflows/.github/workflows/checks.yaml@main
+    with:
+      # Seed the default-branch pre-commit cache without running checks that only apply to PRs.
+      base-branch: main
+      enable_check_size: false
+      enable_check_generated_files: false
+      enable_check_pr_job_dependencies: false
+      enable_check_version_files_changed: false
+      enable_check_version_against_tag: false
   build-details:
     uses: rapidsai/shared-workflows/.github/workflows/compute-build-details.yaml@main
   cpp-build:

.github/workflows/pr.yaml

@@ -434,6 +434,9 @@ jobs:
       contents: read
     with:
       enable_check_generated_files: false
+      # Main builds seed this cache. PR merge refs can restore it but should not create
+      # duplicate archives that are inaccessible to other PRs.
+      pre-commit-cache-read-only: true
   check-c-abi:
     permissions:
       actions: read

c/src/neighbors/mg_cagra.cpp

@@ -325,12 +325,12 @@ void* _mg_deserialize(cuvsResources_t res, const char* filename, mg_cagra_datase
   if (layout == mg_cagra_dataset_layout::device_padded) {
     using padded_ann_t = cuvs::neighbors::cagra::device_padded_index<T, uint32_t>;
     auto* mg_index = new mg_cagra_index_t<T, padded_ann_t>(*res_ptr, cuvs::neighbors::REPLICATED);
-    cuvs::neighbors::cagra::deserialize<T, uint32_t>(*res_ptr, std::string(filename), mg_index);
+    cuvs::neighbors::cagra::deserialize(*res_ptr, std::string(filename), mg_index);
     return make_mg_cagra_box<T, padded_ann_t>(mg_index, mg_cagra_dataset_layout::device_padded);
   }
   using standard_ann_t = cuvs::neighbors::cagra::device_standard_index<T, uint32_t>;
   auto* mg_index       = new mg_cagra_index_t<T, standard_ann_t>(*res_ptr, cuvs::neighbors::REPLICATED);
-  cuvs::neighbors::cagra::deserialize<T, uint32_t>(*res_ptr, std::string(filename), mg_index);
+  cuvs::neighbors::cagra::deserialize(*res_ptr, std::string(filename), mg_index);
   return make_mg_cagra_box<T, standard_ann_t>(mg_index, mg_cagra_dataset_layout::device_standard);
 }
 
@@ -341,12 +341,12 @@ void* _mg_distribute(cuvsResources_t res, const char* filename, mg_cagra_dataset
   if (layout == mg_cagra_dataset_layout::device_padded) {
     using padded_ann_t = cuvs::neighbors::cagra::device_padded_index<T, uint32_t>;
     auto* mg_index = new mg_cagra_index_t<T, padded_ann_t>(*res_ptr, cuvs::neighbors::REPLICATED);
-    cuvs::neighbors::cagra::distribute<T, uint32_t>(*res_ptr, std::string(filename), mg_index);
+    cuvs::neighbors::cagra::distribute(*res_ptr, std::string(filename), mg_index);
     return make_mg_cagra_box<T, padded_ann_t>(mg_index, mg_cagra_dataset_layout::device_padded);
   }
   using standard_ann_t = cuvs::neighbors::cagra::device_standard_index<T, uint32_t>;
   auto* mg_index       = new mg_cagra_index_t<T, standard_ann_t>(*res_ptr, cuvs::neighbors::REPLICATED);
-  cuvs::neighbors::cagra::distribute<T, uint32_t>(*res_ptr, std::string(filename), mg_index);
+  cuvs::neighbors::cagra::distribute(*res_ptr, std::string(filename), mg_index);
   return make_mg_cagra_box<T, standard_ann_t>(mg_index, mg_cagra_dataset_layout::device_standard);
 }
 

cpp/bench/ann/src/cuvs/cuvs_mg_cagra_wrapper.h

@@ -151,7 +151,7 @@ void cuvs_mg_cagra<T, IdxT>::load(const std::string& file)
   index_ = std::make_shared<
     cuvs::neighbors::mg_index<cuvs::neighbors::cagra::device_padded_index<T, IdxT>, T, IdxT>>(
     clique_, index_params_.mode);
-  cuvs::neighbors::cagra::deserialize<T, IdxT>(clique_, file, index_.get());
+  cuvs::neighbors::cagra::deserialize(clique_, file, index_.get());
   query_dim_ = static_cast<uint32_t>(dim_);
   if (!index_->ann_interfaces_.empty() && index_->ann_interfaces_[0].index_.has_value()) {
     query_dim_ = static_cast<uint32_t>(index_->ann_interfaces_[0].index_->dim());

cpp/include/cuvs/neighbors/cagra.hpp

@@ -3494,9 +3494,9 @@ struct merge_params {
 
 /** @brief Merge multiple physical CAGRA indices into one.
  *
- * The overload without `merge_params` uses `merge_algo::AUTO`. AUTO runs Fastener only after a
- * non-mutating preflight validates every input and option; otherwise it calls the existing rebuild
- * implementation. REBUILD always preserves every input dataset.
+ * `merge_params` selects the algorithm and defaults to `merge_algo::AUTO`. AUTO runs Fastener only
+ * after a non-mutating preflight validates every input and option; otherwise it calls the existing
+ * rebuild implementation. REBUILD always preserves every input dataset.
  *
  * Fastener supports unfiltered, uncompressed `float`, `half`, `int8_t`, and `uint8_t`
  * indices using L2Expanded and `uint32_t` graph IDs. Fastener uses `root_fanout` at the first
@@ -3510,40 +3510,91 @@ struct merge_params {
  * which holds only a view of it.
  *
  * @note This API only supports physical merge (`merge_strategy = MERGE_STRATEGY_PHYSICAL`).
- * All input indices must use the same `DatasetViewT` (dense padded or standard device views),
- * and must share one row stride.
+ * All input indices must use the same dataset layout (device padded or device standard), and must
+ * share one row stride.
  *
  * @param[in] res RAFT resources used for the merge.
  * @param[in] params Parameters for the returned CAGRA index.
  * @param[in] indices CAGRA indices to merge.
  * @param[in] merged_dataset Caller-owned storage for the consolidated dataset. Must have one row
  * per retained vector after applying `row_filter`, and the same dimension and stride as the input
  * datasets.
+ * @param[in] merge_params Parameters for the merge, including the algorithm selection.
  * @param[in] row_filter Optional row filter. Any filter selects rebuild in AUTO and is rejected by
  * explicit FASTENER.
  * @return The merged physical CAGRA index.
  */
-template <typename T, typename IdxT, cuvs::neighbors::ann_dataset_view DatasetViewT>
 auto merge(raft::resources const& res,
            const cuvs::neighbors::cagra::index_params& params,
-           std::vector<cuvs::neighbors::cagra::index<T, IdxT, DatasetViewT>*>& indices,
-           DatasetViewT merged_dataset,
+           std::vector<cuvs::neighbors::cagra::device_padded_index<float, uint32_t>*>& indices,
+           cuvs::neighbors::device_padded_dataset_view<float, int64_t> merged_dataset,
+           const cuvs::neighbors::cagra::merge_params& merge_params = {},
            const cuvs::neighbors::filtering::base_filter& row_filter =
              cuvs::neighbors::filtering::none_sample_filter{})
-  -> cuvs::neighbors::cagra::index<T, IdxT, DatasetViewT>;
+  -> cuvs::neighbors::cagra::device_padded_index<float, uint32_t>;
+
+auto merge(raft::resources const& res,
+           const cuvs::neighbors::cagra::index_params& params,
+           std::vector<cuvs::neighbors::cagra::device_standard_index<float, uint32_t>*>& indices,
+           cuvs::neighbors::device_standard_dataset_view<float, int64_t> merged_dataset,
+           const cuvs::neighbors::cagra::merge_params& merge_params = {},
+           const cuvs::neighbors::filtering::base_filter& row_filter =
+             cuvs::neighbors::filtering::none_sample_filter{})
+  -> cuvs::neighbors::cagra::device_standard_index<float, uint32_t>;
+
+auto merge(raft::resources const& res,
+           const cuvs::neighbors::cagra::index_params& params,
+           std::vector<cuvs::neighbors::cagra::device_padded_index<half, uint32_t>*>& indices,
+           cuvs::neighbors::device_padded_dataset_view<half, int64_t> merged_dataset,
+           const cuvs::neighbors::cagra::merge_params& merge_params = {},
+           const cuvs::neighbors::filtering::base_filter& row_filter =
+             cuvs::neighbors::filtering::none_sample_filter{})
+  -> cuvs::neighbors::cagra::device_padded_index<half, uint32_t>;
+
+auto merge(raft::resources const& res,
+           const cuvs::neighbors::cagra::index_params& params,
+           std::vector<cuvs::neighbors::cagra::device_standard_index<half, uint32_t>*>& indices,
+           cuvs::neighbors::device_standard_dataset_view<half, int64_t> merged_dataset,
+           const cuvs::neighbors::cagra::merge_params& merge_params = {},
+           const cuvs::neighbors::filtering::base_filter& row_filter =
+             cuvs::neighbors::filtering::none_sample_filter{})
+  -> cuvs::neighbors::cagra::device_standard_index<half, uint32_t>;
+
+auto merge(raft::resources const& res,
+           const cuvs::neighbors::cagra::index_params& params,
+           std::vector<cuvs::neighbors::cagra::device_padded_index<int8_t, uint32_t>*>& indices,
+           cuvs::neighbors::device_padded_dataset_view<int8_t, int64_t> merged_dataset,
+           const cuvs::neighbors::cagra::merge_params& merge_params = {},
+           const cuvs::neighbors::filtering::base_filter& row_filter =
+             cuvs::neighbors::filtering::none_sample_filter{})
+  -> cuvs::neighbors::cagra::device_padded_index<int8_t, uint32_t>;
 
-/** @copydoc merge
- * @param[in] merge_params Parameters for the merge, including the algorithm selection.
- */
-template <typename T, typename IdxT, cuvs::neighbors::ann_dataset_view DatasetViewT>
 auto merge(raft::resources const& res,
            const cuvs::neighbors::cagra::index_params& params,
-           std::vector<cuvs::neighbors::cagra::index<T, IdxT, DatasetViewT>*>& indices,
-           DatasetViewT merged_dataset,
-           const cuvs::neighbors::cagra::merge_params& merge_params,
+           std::vector<cuvs::neighbors::cagra::device_standard_index<int8_t, uint32_t>*>& indices,
+           cuvs::neighbors::device_standard_dataset_view<int8_t, int64_t> merged_dataset,
+           const cuvs::neighbors::cagra::merge_params& merge_params = {},
            const cuvs::neighbors::filtering::base_filter& row_filter =
              cuvs::neighbors::filtering::none_sample_filter{})
-  -> cuvs::neighbors::cagra::index<T, IdxT, DatasetViewT>;
+  -> cuvs::neighbors::cagra::device_standard_index<int8_t, uint32_t>;
+
+auto merge(raft::resources const& res,
+           const cuvs::neighbors::cagra::index_params& params,
+           std::vector<cuvs::neighbors::cagra::device_padded_index<uint8_t, uint32_t>*>& indices,
+           cuvs::neighbors::device_padded_dataset_view<uint8_t, int64_t> merged_dataset,
+           const cuvs::neighbors::cagra::merge_params& merge_params = {},
+           const cuvs::neighbors::filtering::base_filter& row_filter =
+             cuvs::neighbors::filtering::none_sample_filter{})
+  -> cuvs::neighbors::cagra::device_padded_index<uint8_t, uint32_t>;
+
+auto merge(raft::resources const& res,
+           const cuvs::neighbors::cagra::index_params& params,
+           std::vector<cuvs::neighbors::cagra::device_standard_index<uint8_t, uint32_t>*>& indices,
+           cuvs::neighbors::device_standard_dataset_view<uint8_t, int64_t> merged_dataset,
+           const cuvs::neighbors::cagra::merge_params& merge_params = {},
+           const cuvs::neighbors::filtering::base_filter& row_filter =
+             cuvs::neighbors::filtering::none_sample_filter{})
+  -> cuvs::neighbors::cagra::device_standard_index<uint8_t, uint32_t>;
 
 /**
  * @}
@@ -4388,23 +4439,56 @@ void serialize(const raft::resources& clique,
  * cuvs::neighbors::cagra::serialize(clique, index, filename);
  * cuvs::neighbors::mg_index<cagra::device_standard_index<float, uint32_t>, float, uint32_t>
  *   new_index(clique, REPLICATED);
- * cuvs::neighbors::cagra::deserialize<float, uint32_t>(clique, filename, &new_index);
+ * cuvs::neighbors::cagra::deserialize(clique, filename, &new_index);
  *
  * @endcode
  *
  * @param[in] clique a `raft::resources` object specifying the NCCL clique configuration
  * @param[in] filename path to the file to be deserialized
  *
  */
-template <typename T, typename IdxT>
-void deserialize(const raft::resources& clique,
-                 const std::string& filename,
-                 cuvs::neighbors::mg_index<cagra::device_standard_index<T, IdxT>, T, IdxT>* index);
+void deserialize(
+  const raft::resources& clique,
+  const std::string& filename,
+  cuvs::neighbors::mg_index<cagra::device_padded_index<float, uint32_t>, float, uint32_t>* index);
 
-template <typename T, typename IdxT>
-void deserialize(const raft::resources& clique,
-                 const std::string& filename,
-                 cuvs::neighbors::mg_index<cagra::device_padded_index<T, IdxT>, T, IdxT>* index);
+void deserialize(
+  const raft::resources& clique,
+  const std::string& filename,
+  cuvs::neighbors::mg_index<cagra::device_standard_index<float, uint32_t>, float, uint32_t>* index);
+
+void deserialize(
+  const raft::resources& clique,
+  const std::string& filename,
+  cuvs::neighbors::mg_index<cagra::device_padded_index<half, uint32_t>, half, uint32_t>* index);
+
+void deserialize(
+  const raft::resources& clique,
+  const std::string& filename,
+  cuvs::neighbors::mg_index<cagra::device_standard_index<half, uint32_t>, half, uint32_t>* index);
+
+void deserialize(
+  const raft::resources& clique,
+  const std::string& filename,
+  cuvs::neighbors::mg_index<cagra::device_padded_index<int8_t, uint32_t>, int8_t, uint32_t>* index);
+
+void deserialize(
+  const raft::resources& clique,
+  const std::string& filename,
+  cuvs::neighbors::mg_index<cagra::device_standard_index<int8_t, uint32_t>, int8_t, uint32_t>*
+    index);
+
+void deserialize(
+  const raft::resources& clique,
+  const std::string& filename,
+  cuvs::neighbors::mg_index<cagra::device_padded_index<uint8_t, uint32_t>, uint8_t, uint32_t>*
+    index);
+
+void deserialize(
+  const raft::resources& clique,
+  const std::string& filename,
+  cuvs::neighbors::mg_index<cagra::device_standard_index<uint8_t, uint32_t>, uint8_t, uint32_t>*
+    index);
 
 /// \defgroup mg_cpp_distribute ANN MG local index distribution
 
@@ -4422,23 +4506,56 @@ void deserialize(const raft::resources& clique,
  * cuvs::neighbors::cagra::serialize(clique, filename, index);
  * cuvs::neighbors::mg_index<cagra::device_standard_index<float, uint32_t>, float, uint32_t>
  *   distributed_index(clique, REPLICATED);
- * cuvs::neighbors::cagra::distribute<float, uint32_t>(clique, filename, &distributed_index);
+ * cuvs::neighbors::cagra::distribute(clique, filename, &distributed_index);
  *
  * @endcode
  *
  * @param[in] clique a `raft::resources` object specifying the NCCL clique configuration
  * @param[in] filename path to the file to be deserialized : a local index
  *
  */
-template <typename T, typename IdxT>
-void distribute(const raft::resources& clique,
-                const std::string& filename,
-                cuvs::neighbors::mg_index<cagra::device_standard_index<T, IdxT>, T, IdxT>* index);
+void distribute(
+  const raft::resources& clique,
+  const std::string& filename,
+  cuvs::neighbors::mg_index<cagra::device_padded_index<float, uint32_t>, float, uint32_t>* index);
 
-template <typename T, typename IdxT>
-void distribute(const raft::resources& clique,
-                const std::string& filename,
-                cuvs::neighbors::mg_index<cagra::device_padded_index<T, IdxT>, T, IdxT>* index);
+void distribute(
+  const raft::resources& clique,
+  const std::string& filename,
+  cuvs::neighbors::mg_index<cagra::device_standard_index<float, uint32_t>, float, uint32_t>* index);
+
+void distribute(
+  const raft::resources& clique,
+  const std::string& filename,
+  cuvs::neighbors::mg_index<cagra::device_padded_index<half, uint32_t>, half, uint32_t>* index);
+
+void distribute(
+  const raft::resources& clique,
+  const std::string& filename,
+  cuvs::neighbors::mg_index<cagra::device_standard_index<half, uint32_t>, half, uint32_t>* index);
+
+void distribute(
+  const raft::resources& clique,
+  const std::string& filename,
+  cuvs::neighbors::mg_index<cagra::device_padded_index<int8_t, uint32_t>, int8_t, uint32_t>* index);
+
+void distribute(
+  const raft::resources& clique,
+  const std::string& filename,
+  cuvs::neighbors::mg_index<cagra::device_standard_index<int8_t, uint32_t>, int8_t, uint32_t>*
+    index);
+
+void distribute(
+  const raft::resources& clique,
+  const std::string& filename,
+  cuvs::neighbors::mg_index<cagra::device_padded_index<uint8_t, uint32_t>, uint8_t, uint32_t>*
+    index);
+
+void distribute(
+  const raft::resources& clique,
+  const std::string& filename,
+  cuvs::neighbors::mg_index<cagra::device_standard_index<uint8_t, uint32_t>, uint8_t, uint32_t>*
+    index);
 
 /**
  * @brief Build a kNN graph using IVF-PQ.