Repository: rapidsai/cuvs
Pull request #888: Expose ivf-flat centers to python/c
Author: login: benfred
Labels: name: improvement
description: Improves an existing functionality
name: non-breaking
description: Introduces a non-breaking change
name: cpp
name: Python
Created: 2025-05-13T00:17:57Z | Merged: 2025-05-27T23:01:47Z

Similar to #881 - also expose centers for ivf-flat as well as ivf-pq




Diff:
diff --git a/cpp/include/cuvs/neighbors/ivf_flat.h b/cpp/include/cuvs/neighbors/ivf_flat.h
index c46a9b2cdb..e274af93b9 100644
--- a/cpp/include/cuvs/neighbors/ivf_flat.h
+++ b/cpp/include/cuvs/neighbors/ivf_flat.h
@@ -164,6 +164,25 @@ cuvsError_t cuvsIvfFlatIndexCreate(cuvsIvfFlatIndex_t* index);
  * @param[in] index cuvsIvfFlatIndex_t to de-allocate
  */
 cuvsError_t cuvsIvfFlatIndexDestroy(cuvsIvfFlatIndex_t index);
+
+/** Get the number of clusters/inverted lists */
+uint32_t cuvsIvfFlatIndexGetNLists(cuvsIvfFlatIndex_t index);
+
+/** Get the dimensionality of the data */
+uint32_t cuvsIvfFlatIndexGetDim(cuvsIvfFlatIndex_t index);
+
+/**
+ * @brief Get the cluster centers corresponding to the lists [n_lists, dim]
+ *
+ * @param[in] res cuvsResources_t opaque C handle
+ * @param[in] index cuvsIvfFlatIndex_t Built Ivf-Flat Index
+ * @param[out] centers Preallocated array on host or device memory to store output, [n_lists, dim]
+ * @return cuvsError_t
+ */
+cuvsError_t cuvsIvfFlatIndexGetCenters(cuvsResources_t res,
+                                       cuvsIvfFlatIndex_t index,
+                                       DLManagedTensor* centers);
+
 /**
  * @}
  */
diff --git a/cpp/include/cuvs/neighbors/ivf_pq.h b/cpp/include/cuvs/neighbors/ivf_pq.h
index b00418ecb8..40b9e9b839 100644
--- a/cpp/include/cuvs/neighbors/ivf_pq.h
+++ b/cpp/include/cuvs/neighbors/ivf_pq.h
@@ -264,15 +264,16 @@ cuvsError_t cuvsIvfPqIndexDestroy(cuvsIvfPqIndex_t index);
 /** Get the number of clusters/inverted lists */
 uint32_t cuvsIvfPqIndexGetNLists(cuvsIvfPqIndex_t index);
 
-/** Get the dimensionality of the cluster centers */
-uint32_t cuvsIvfPqIndexGetDimExt(cuvsIvfPqIndex_t index);
+/** Get the dimensionality */
+uint32_t cuvsIvfPqIndexGetDim(cuvsIvfPqIndex_t index);
 
 /**
  * @brief Get the cluster centers corresponding to the lists in the original space
  *
  * @param[in] res cuvsResources_t opaque C handle
- * @param[in] index cuvsIvfPqIndex_t Built NN-Descent index
- * @param[out] centers Preallocated array on host memory to store output, [n_lists, dim_ext]
+ * @param[in] index cuvsIvfPqIndex_t Built Ivf-Pq index
+ * @param[out] centers Preallocated array on host or device memory to store output,
+ * dimensions [n_lists, dim]
  * @return cuvsError_t
  */
 cuvsError_t cuvsIvfPqIndexGetCenters(cuvsResources_t res,
diff --git a/cpp/include/cuvs/neighbors/ivf_pq.hpp b/cpp/include/cuvs/neighbors/ivf_pq.hpp
index 7c2a3e9f85..2db1f33eb2 100644
--- a/cpp/include/cuvs/neighbors/ivf_pq.hpp
+++ b/cpp/include/cuvs/neighbors/ivf_pq.hpp
@@ -2861,6 +2861,10 @@ void extract_centers(raft::resources const& res,
                      const index<int64_t>& index,
                      raft::device_matrix_view<float, uint32_t, raft::row_major> cluster_centers);
 
+/** @copydoc extract_centers */
+void extract_centers(raft::resources const& res,
+                     const index<int64_t>& index,
+                     raft::host_matrix_view<float, uint32_t, raft::row_major> cluster_centers);
 /**
  * @brief Helper exposing the re-computation of list sizes and related arrays if IVF lists have been
  * modified externally.
diff --git a/cpp/src/neighbors/ivf_flat_c.cpp b/cpp/src/neighbors/ivf_flat_c.cpp
index 4f930521a6..df02ea17ca 100644
--- a/cpp/src/neighbors/ivf_flat_c.cpp
+++ b/cpp/src/neighbors/ivf_flat_c.cpp
@@ -22,6 +22,7 @@
 #include <raft/core/mdspan_types.hpp>
 #include <raft/core/resources.hpp>
 #include <raft/core/serialize.hpp>
+#include <raft/util/cudart_utils.hpp>
 
 #include <cuvs/core/c_api.h>
 #include <cuvs/core/exceptions.hpp>
@@ -139,6 +140,35 @@ void _extend(cuvsResources_t res,
 
   cuvs::neighbors::ivf_flat::extend(*res_ptr, vectors_mds, indices_mds, index_ptr);
 }
+
+template <typename output_mdspan_type, typename T, typename IdxT>
+void _get_centers(cuvsResources_t res, cuvsIvfFlatIndex index, DLManagedTensor* centers)
+{
+  auto res_ptr   = reinterpret_cast<raft::resources*>(res);
+  auto index_ptr = reinterpret_cast<cuvs::neighbors::ivf_flat::index<T, IdxT>*>(index.addr);
+  auto dst       = cuvs::core::from_dlpack<output_mdspan_type>(centers);
+  auto src       = index_ptr->centers();
+
+  RAFT_EXPECTS(src.extent(0) == dst.extent(0), "Output centers has incorrect number of rows");
+  RAFT_EXPECTS(src.extent(1) == dst.extent(1), "Output centers has incorrect number of cols");
+
+  raft::copy(dst.data_handle(),
+             src.data_handle(),
+             dst.extent(0) * dst.extent(1),
+             raft::resource::get_cuda_stream(*res_ptr));
+}
+
+template <typename T, typename IdxT>
+void get_centers(cuvsResources_t res, cuvsIvfFlatIndex index, DLManagedTensor* centers)
+{
+  if (cuvs::core::is_dlpack_device_compatible(centers->dl_tensor)) {
+    using output_mdspan_type = raft::device_matrix_view<float, int64_t, raft::row_major>;
+    _get_centers<output_mdspan_type, T, IdxT>(res, index, centers);
+  } else {
+    using output_mdspan_type = raft::host_matrix_view<float, int64_t, raft::row_major>;
+    _get_centers<output_mdspan_type, T, IdxT>(res, index, centers);
+  }
+}
 }  // namespace
 
 extern "C" cuvsError_t cuvsIvfFlatIndexCreate(cuvsIvfFlatIndex_t* index)
@@ -351,3 +381,68 @@ extern "C" cuvsError_t cuvsIvfFlatExtend(cuvsResources_t res,
     }
   });
 }
+
+extern "C" uint32_t cuvsIvfFlatIndexGetNLists(cuvsIvfFlatIndex_t index)
+{
+  if (index->dtype.code == kDLFloat && index->dtype.bits == 32) {
+    auto index_ptr =
+      reinterpret_cast<cuvs::neighbors::ivf_flat::index<float, int64_t>*>(index->addr);
+    return index_ptr->n_lists();
+  } else if (index->dtype.code == kDLFloat && index->dtype.bits == 16) {
+    auto index_ptr =
+      reinterpret_cast<cuvs::neighbors::ivf_flat::index<half, int64_t>*>(index->addr);
+    return index_ptr->n_lists();
+  } else if (index->dtype.code == kDLInt && index->dtype.bits == 8) {
+    auto index_ptr =
+      reinterpret_cast<cuvs::neighbors::ivf_flat::index<int8_t, int64_t>*>(index->addr);
+    return index_ptr->n_lists();
+  } else if (index->dtype.code == kDLUInt && index->dtype.bits == 8) {
+    auto index_ptr =
+      reinterpret_cast<cuvs::neighbors::ivf_flat::index<uint8_t, int64_t>*>(index->addr);
+    return index_ptr->n_lists();
+  } else {
+    return 0;
+  }
+}
+
+extern "C" uint32_t cuvsIvfFlatIndexGetDim(cuvsIvfFlatIndex_t index)
+{
+  if (index->dtype.code == kDLFloat && index->dtype.bits == 32) {
+    auto index_ptr =
+      reinterpret_cast<cuvs::neighbors::ivf_flat::index<float, int64_t>*>(index->addr);
+    return index_ptr->dim();
+  } else if (index->dtype.code == kDLFloat && index->dtype.bits == 16) {
+    auto index_ptr =
+      reinterpret_cast<cuvs::neighbors::ivf_flat::index<half, int64_t>*>(index->addr);
+    return index_ptr->dim();
+  } else if (index->dtype.code == kDLInt && index->dtype.bits == 8) {
+    auto index_ptr =
+      reinterpret_cast<cuvs::neighbors::ivf_flat::index<int8_t, int64_t>*>(index->addr);
+    return index_ptr->dim();
+  } else if (index->dtype.code == kDLUInt && index->dtype.bits == 8) {
+    auto index_ptr =
+      reinterpret_cast<cuvs::neighbors::ivf_flat::index<uint8_t, int64_t>*>(index->addr);
+    return index_ptr->dim();
+  } else {
+    return 0;
+  }
+}
+
+extern "C" cuvsError_t cuvsIvfFlatIndexGetCenters(cuvsResources_t res,
+                                                  cuvsIvfFlatIndex_t index,
+                                                  DLManagedTensor* centers)
+{
+  return cuvs::core::translate_exceptions([=] {
+    if (index->dtype.code == kDLFloat && index->dtype.bits == 32) {
+      get_centers<float, int64_t>(res, *index, centers);
+    } else if (index->dtype.code == kDLFloat && index->dtype.bits == 16) {
+      get_centers<half, int64_t>(res, *index, centers);
+    } else if (index->dtype.code == kDLInt && index->dtype.bits == 8) {
+      get_centers<int8_t, int64_t>(res, *index, centers);
+    } else if (index->dtype.code == kDLUInt && index->dtype.bits == 8) {
+      get_centers<uint8_t, int64_t>(res, *index, centers);
+    } else {
+      RAFT_FAIL("Unsupported index dtype: %d and bits: %d", index->dtype.code, index->dtype.bits);
+    }
+  });
+}
diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh b/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh
index 2ee0335f6e..aaeb2cdd5e 100644
--- a/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh
+++ b/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh
@@ -1904,4 +1904,26 @@ void extend(
                   new_indices.has_value() ? new_indices.value().data_handle() : nullptr,
                   n_rows);
 }
+
+template <typename output_mdspan_type>
+inline void extract_centers(raft::resources const& res,
+                            const cuvs::neighbors::ivf_pq::index<int64_t>& index,
+                            output_mdspan_type cluster_centers)
+{
+  RAFT_EXPECTS(cluster_centers.extent(0) == index.n_lists(),
+               "Number of rows in the output buffer for cluster centers must be equal to the "
+               "number of IVF lists");
+  RAFT_EXPECTS(
+    cluster_centers.extent(1) == index.dim(),
+    "Number of columns in the output buffer for cluster centers and index dim are different");
+  auto stream = raft::resource::get_cuda_stream(res);
+  RAFT_CUDA_TRY(cudaMemcpy2DAsync(cluster_centers.data_handle(),
+                                  sizeof(float) * index.dim(),
+                                  index.centers().data_handle(),
+                                  sizeof(float) * index.dim_ext(),
+                                  sizeof(float) * index.dim(),
+                                  index.n_lists(),
+                                  cudaMemcpyDefault,
+                                  stream));
+}
 }  // namespace cuvs::neighbors::ivf_pq::detail
diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_build_common.cu b/cpp/src/neighbors/ivf_pq/ivf_pq_build_common.cu
index b34611c1d6..dc54964275 100644
--- a/cpp/src/neighbors/ivf_pq/ivf_pq_build_common.cu
+++ b/cpp/src/neighbors/ivf_pq/ivf_pq_build_common.cu
@@ -294,21 +294,14 @@ void extract_centers(raft::resources const& res,
                      const cuvs::neighbors::ivf_pq::index<int64_t>& index,
                      raft::device_matrix_view<float, uint32_t, raft::row_major> cluster_centers)
 {
-  RAFT_EXPECTS(cluster_centers.extent(0) == index.n_lists(),
-               "Number of rows in the output buffer for cluster centers must be equal to the "
-               "number of IVF lists");
-  RAFT_EXPECTS(
-    cluster_centers.extent(1) == index.dim(),
-    "Number of columns in the output buffer for cluster centers and index dim are different");
-  auto stream = raft::resource::get_cuda_stream(res);
-  RAFT_CUDA_TRY(cudaMemcpy2DAsync(cluster_centers.data_handle(),
-                                  sizeof(float) * index.dim(),
-                                  index.centers().data_handle(),
-                                  sizeof(float) * index.dim_ext(),
-                                  sizeof(float) * index.dim(),
-                                  index.n_lists(),
-                                  cudaMemcpyDefault,
-                                  stream));
+  detail::extract_centers(res, index, cluster_centers);
+}
+
+void extract_centers(raft::resources const& res,
+                     const cuvs::neighbors::ivf_pq::index<int64_t>& index,
+                     raft::host_matrix_view<float, uint32_t, raft::row_major> cluster_centers)
+{
+  detail::extract_centers(res, index, cluster_centers);
 }
 
 void recompute_internal_state(const raft::resources& res, index<int64_t>* index)
diff --git a/cpp/src/neighbors/ivf_pq_c.cpp b/cpp/src/neighbors/ivf_pq_c.cpp
index 0a3db6724a..2412bcbcc0 100644
--- a/cpp/src/neighbors/ivf_pq_c.cpp
+++ b/cpp/src/neighbors/ivf_pq_c.cpp
@@ -150,16 +150,8 @@ void _get_centers(cuvsResources_t res, cuvsIvfPqIndex index, DLManagedTensor* ce
   auto res_ptr   = reinterpret_cast<raft::resources*>(res);
   auto index_ptr = reinterpret_cast<cuvs::neighbors::ivf_pq::index<IdxT>*>(index.addr);
   auto dst       = cuvs::core::from_dlpack<output_mdspan_type>(centers);
-  auto src       = index_ptr->centers();
 
-  RAFT_EXPECTS(src.extent(0) == dst.extent(0), "Output centers has incorrect number of rows");
-  RAFT_EXPECTS(src.extent(1) == dst.extent(1), "Output centers has incorrect number of cols");
-
-  cudaMemcpyAsync(dst.data_handle(),
-                  src.data_handle(),
-                  dst.extent(0) * dst.extent(1) * sizeof(float),
-                  cudaMemcpyDefault,
-                  raft::resource::get_cuda_stream(*res_ptr));
+  cuvs::neighbors::ivf_pq::helpers::extract_centers(*res_ptr, *index_ptr, dst);
 }
 }  // namespace
 
@@ -337,10 +329,10 @@ extern "C" uint32_t cuvsIvfPqIndexGetNLists(cuvsIvfPqIndex_t index)
   return index_ptr->n_lists();
 }
 
-extern "C" uint32_t cuvsIvfPqIndexGetDimExt(cuvsIvfPqIndex_t index)
+extern "C" uint32_t cuvsIvfPqIndexGetDim(cuvsIvfPqIndex_t index)
 {
   auto index_ptr = reinterpret_cast<cuvs::neighbors::ivf_pq::index<int64_t>*>(index->addr);
-  return index_ptr->dim_ext();
+  return index_ptr->dim();
 }
 
 extern "C" cuvsError_t cuvsIvfPqIndexGetCenters(cuvsResources_t res,
diff --git a/cpp/src/neighbors/nn_descent_c.cpp b/cpp/src/neighbors/nn_descent_c.cpp
index a3cc760ef5..ddebb11e73 100644
--- a/cpp/src/neighbors/nn_descent_c.cpp
+++ b/cpp/src/neighbors/nn_descent_c.cpp
@@ -23,6 +23,7 @@
 #include <raft/core/resource/cuda_stream.hpp>
 #include <raft/core/resources.hpp>
 #include <raft/core/serialize.hpp>
+#include <raft/util/cudart_utils.hpp>
 
 #include <cuvs/core/c_api.h>
 #include <cuvs/core/exceptions.hpp>
@@ -85,11 +86,10 @@ void _get_graph(cuvsResources_t res, cuvsNNDescentIndex_t index, DLManagedTensor
     RAFT_EXPECTS(src.extent(0) == dst.extent(0), "Output graph has incorrect number of rows");
     RAFT_EXPECTS(src.extent(1) == dst.extent(1), "Output graph has incorrect number of cols");
 
-    cudaMemcpyAsync(dst.data_handle(),
-                    src.data_handle(),
-                    dst.extent(0) * dst.extent(1) * sizeof(uint32_t),
-                    cudaMemcpyDefault,
-                    raft::resource::get_cuda_stream(*res_ptr));
+    raft::copy(dst.data_handle(),
+               src.data_handle(),
+               dst.extent(0) * dst.extent(1),
+               raft::resource::get_cuda_stream(*res_ptr));
   } else {
     RAFT_FAIL("Unsupported nn-descent index dtype: %d and bits: %d", dtype.code, dtype.bits);
   }
diff --git a/python/cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pxd b/python/cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pxd
index f2bd6a9b18..b97faa1ad9 100644
--- a/python/cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pxd
+++ b/python/cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pxd
@@ -62,6 +62,14 @@ cdef extern from "cuvs/neighbors/ivf_flat.h" nogil:
 
     cuvsError_t cuvsIvfFlatIndexDestroy(cuvsIvfFlatIndex_t index)
 
+    uint32_t cuvsIvfFlatIndexGetNLists(cuvsIvfFlatIndex_t index)
+
+    uint32_t cuvsIvfFlatIndexGetDim(cuvsIvfFlatIndex_t index)
+
+    cuvsError_t cuvsIvfFlatIndexGetCenters(cuvsResources_t res,
+                                           cuvsIvfFlatIndex_t index,
+                                           DLManagedTensor * centers)
+
     cuvsError_t cuvsIvfFlatBuild(cuvsResources_t res,
                                  cuvsIvfFlatIndexParams* params,
                                  DLManagedTensor* dataset,
diff --git a/python/cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pyx b/python/cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pyx
index 658a32c838..4db2854365 100644
--- a/python/cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pyx
+++ b/python/cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pyx
@@ -181,6 +181,35 @@ cdef class Index:
     def __repr__(self):
         return "Index(type=IvfFlat)"
 
+    @property
+    def n_lists(self):
+        """ The number of inverted lists (clusters) """
+        return cuvsIvfFlatIndexGetNLists(self.index)
+
+    @property
+    def dim(self):
+        """ dimensionality of the cluster centers """
+        return cuvsIvfFlatIndexGetDim(self.index)
+
+    @property
+    def centers(self):
+        """ Get the cluster centers corresponding to the lists in the
+        original space """
+        return self._get_centers()
+
+    @auto_sync_resources
+    def _get_centers(self, resources=None):
+        if not self.trained:
+            raise ValueError("Index needs to be built before getting centers")
+
+        cdef cuvsResources_t res = <cuvsResources_t>resources.get_c_obj()
+
+        output = np.empty((self.n_lists, self.dim), dtype='float32')
+        ai = wrap_array(output)
+        cdef cydlpack.DLManagedTensor* output_dlpack = cydlpack.dlpack_c(ai)
+        check_cuvs(cuvsIvfFlatIndexGetCenters(res, self.index, output_dlpack))
+        return output
+
 
 @auto_sync_resources
 def build(IndexParams index_params, dataset, resources=None):
diff --git a/python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pxd b/python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pxd
index 1ed2091644..635bb970d5 100644
--- a/python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pxd
+++ b/python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pxd
@@ -83,7 +83,7 @@ cdef extern from "cuvs/neighbors/ivf_pq.h" nogil:
 
     uint32_t cuvsIvfPqIndexGetNLists(cuvsIvfPqIndex_t index)
 
-    uint32_t cuvsIvfPqIndexGetDimExt(cuvsIvfPqIndex_t index)
+    uint32_t cuvsIvfPqIndexGetDim(cuvsIvfPqIndex_t index)
 
     cuvsError_t cuvsIvfPqIndexGetCenters(cuvsResources_t res,
                                          cuvsIvfPqIndex_t index,
diff --git a/python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pyx b/python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pyx
index 374279b31b..678ca2705c 100644
--- a/python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pyx
+++ b/python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pyx
@@ -244,9 +244,9 @@ cdef class Index:
         return cuvsIvfPqIndexGetNLists(self.index)
 
     @property
-    def dim_ext(self):
+    def dim(self):
         """ dimensionality of the cluster centers """
-        return cuvsIvfPqIndexGetDimExt(self.index)
+        return cuvsIvfPqIndexGetDim(self.index)
 
     @property
     def centers(self):
@@ -261,7 +261,7 @@ cdef class Index:
 
         cdef cuvsResources_t res = <cuvsResources_t>resources.get_c_obj()
 
-        output = np.empty((self.n_lists, self.dim_ext), dtype='float32')
+        output = np.empty((self.n_lists, self.dim), dtype='float32')
         ai = wrap_array(output)
         cdef cydlpack.DLManagedTensor* output_dlpack = cydlpack.dlpack_c(ai)
         check_cuvs(cuvsIvfPqIndexGetCenters(res, self.index, output_dlpack))
diff --git a/python/cuvs/cuvs/tests/test_ivf_flat.py b/python/cuvs/cuvs/tests/test_ivf_flat.py
index e0cd7b08b1..ea51f3a02c 100644
--- a/python/cuvs/cuvs/tests/test_ivf_flat.py
+++ b/python/cuvs/cuvs/tests/test_ivf_flat.py
@@ -108,6 +108,10 @@ def run_ivf_flat_build_search_test(
     recall = calc_recall(out_idx, skl_idx)
     assert recall > 0.7
 
+    centers = index.centers
+    assert centers.shape[0] == build_params.n_lists
+    assert centers.shape[1] == n_cols
+
 
 @pytest.mark.parametrize("inplace", [True, False])
 @pytest.mark.parametrize("dtype", [np.float32])


Reviews:
user: login: tarang-jain
state: CHANGES_REQUESTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-13T19:07:14Z
user: login: benfred
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-13T19:28:10Z
user: login: tarang-jain
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-13T19:51:45Z
user: login: cjnolet
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-13T20:22:54Z
user: login: benfred
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-13T20:27:01Z
user: login: tarang-jain
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-13T20:44:21Z
user: login: cjnolet
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-13T20:47:47Z
user: login: tarang-jain
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-13T20:56:32Z
user: login: benfred
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-13T20:56:59Z
user: login: tarang-jain
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-13T20:57:43Z
user: login: cjnolet
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-13T20:58:01Z
user: login: benfred
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-13T20:59:11Z
user: login: tarang-jain
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-13T21:08:00Z
user: login: benfred
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-13T22:34:15Z
user: login: tarang-jain
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-14T18:50:54Z
user: login: mythrocks
body: Some trivial nitpicks.  Still coming to terms with the code.

The C++ side looks good to my eye.  +1 non-binding.
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-19T18:35:16Z
user: login: cjnolet
state: APPROVED
author_association: MEMBER
submitted_at: 2025-05-27T22:18:20Z

Review comments:
diff_hunk: @@ -139,6 +139,36 @@ void _extend(cuvsResources_t res,
 
   cuvs::neighbors::ivf_flat::extend(*res_ptr, vectors_mds, indices_mds, index_ptr);
 }
+
+template <typename output_mdspan_type, typename T, typename IdxT>
+void _get_centers(cuvsResources_t res, cuvsIvfFlatIndex index, DLManagedTensor* centers)
+{
+  auto res_ptr   = reinterpret_cast<raft::resources*>(res);
+  auto index_ptr = reinterpret_cast<cuvs::neighbors::ivf_flat::index<T, IdxT>*>(index.addr);
+  auto dst       = cuvs::core::from_dlpack<output_mdspan_type>(centers);
+  auto src       = index_ptr->centers();
+
+  RAFT_EXPECTS(src.extent(0) == dst.extent(0), "Output centers has incorrect number of rows");
+  RAFT_EXPECTS(src.extent(1) == dst.extent(1), "Output centers has incorrect number of cols");
+
+  cudaMemcpyAsync(dst.data_handle(),
path: cpp/src/neighbors/ivf_flat_c.cpp
user: login: tarang-jain
body: Can we use `raft::copy` or `raft::copy_async` in these calls instead of direct calls to `cudaMemcpyAsync`?
created_at: 2025-05-13T18:56:08Z
author_association: CONTRIBUTOR
diff_hunk: @@ -139,6 +139,36 @@ void _extend(cuvsResources_t res,
 
   cuvs::neighbors::ivf_flat::extend(*res_ptr, vectors_mds, indices_mds, index_ptr);
 }
+
+template <typename output_mdspan_type, typename T, typename IdxT>
+void _get_centers(cuvsResources_t res, cuvsIvfFlatIndex index, DLManagedTensor* centers)
+{
+  auto res_ptr   = reinterpret_cast<raft::resources*>(res);
+  auto index_ptr = reinterpret_cast<cuvs::neighbors::ivf_flat::index<T, IdxT>*>(index.addr);
+  auto dst       = cuvs::core::from_dlpack<output_mdspan_type>(centers);
+  auto src       = index_ptr->centers();
+
+  RAFT_EXPECTS(src.extent(0) == dst.extent(0), "Output centers has incorrect number of rows");
+  RAFT_EXPECTS(src.extent(1) == dst.extent(1), "Output centers has incorrect number of cols");
+
+  cudaMemcpyAsync(dst.data_handle(),
path: cpp/src/neighbors/ivf_flat_c.cpp
user: login: benfred
body: sure! I updated in the last commit to use the `raft::copy` from `raft/util/cudart_utils.hpp` .

(Fwiw, I originally tried to use the `copy` functions from `raft/core/copy.hpp` and `raft/core/copy.cuh` but couldn't - since we can't use the '.cuh' version inside the C-api, and the `.hpp` version was complaining about needing a cuda kernel for the D2H copy iirc, which is why I was using cudaMemCpyAsync directly here).
created_at: 2025-05-13T19:28:10Z
author_association: CONTRIBUTOR
diff_hunk: @@ -139,6 +139,36 @@ void _extend(cuvsResources_t res,
 
   cuvs::neighbors::ivf_flat::extend(*res_ptr, vectors_mds, indices_mds, index_ptr);
 }
+
+template <typename output_mdspan_type, typename T, typename IdxT>
+void _get_centers(cuvsResources_t res, cuvsIvfFlatIndex index, DLManagedTensor* centers)
+{
+  auto res_ptr   = reinterpret_cast<raft::resources*>(res);
+  auto index_ptr = reinterpret_cast<cuvs::neighbors::ivf_flat::index<T, IdxT>*>(index.addr);
+  auto dst       = cuvs::core::from_dlpack<output_mdspan_type>(centers);
+  auto src       = index_ptr->centers();
+
+  RAFT_EXPECTS(src.extent(0) == dst.extent(0), "Output centers has incorrect number of rows");
+  RAFT_EXPECTS(src.extent(1) == dst.extent(1), "Output centers has incorrect number of cols");
+
+  cudaMemcpyAsync(dst.data_handle(),
path: cpp/src/neighbors/ivf_flat_c.cpp
user: login: tarang-jain
body: I am seeing that `nn_descent_c.cpp` also uses `cudaMemcpyAsync`. If there is a reason we have avoided the cudart header, we can stick to `cudaMemcpyAsync`. My understanding was that having the `cudaMemcpyAsync` call means that the file would have to be compiled with nvcc anyway so we should be reusing raft functions instead.
(cc @cjnolet)
created_at: 2025-05-13T19:51:45Z
author_association: CONTRIBUTOR
diff_hunk: @@ -139,6 +139,36 @@ void _extend(cuvsResources_t res,
 
   cuvs::neighbors::ivf_flat::extend(*res_ptr, vectors_mds, indices_mds, index_ptr);
 }
+
+template <typename output_mdspan_type, typename T, typename IdxT>
+void _get_centers(cuvsResources_t res, cuvsIvfFlatIndex index, DLManagedTensor* centers)
+{
+  auto res_ptr   = reinterpret_cast<raft::resources*>(res);
+  auto index_ptr = reinterpret_cast<cuvs::neighbors::ivf_flat::index<T, IdxT>*>(index.addr);
+  auto dst       = cuvs::core::from_dlpack<output_mdspan_type>(centers);
+  auto src       = index_ptr->centers();
+
+  RAFT_EXPECTS(src.extent(0) == dst.extent(0), "Output centers has incorrect number of rows");
+  RAFT_EXPECTS(src.extent(1) == dst.extent(1), "Output centers has incorrect number of cols");
+
+  cudaMemcpyAsync(dst.data_handle(),
path: cpp/src/neighbors/ivf_flat_c.cpp
user: login: cjnolet
body: >  My understanding was that having the cudaMemcpyAsync call means that the file would have to be compiled with nvcc anyway 

@tarang-jain `cudaMemCpyXX()` is part of the CUDA runtime API, so it should not require `nvcc` to compile. Only the lower-level device function routines will need `nvcc`. Otherwise, it's just linking against the pre-compiled CUDA routines in the runtime API (kind of similar to what end-users do when they use cuVS C/C++ APIs). 

The reason why `raft::copy` ends up requiring `nvcc` is because there were recently some device functions added to raft to work specifically with mdspan... to be honest, I'd be in favor of separating those out from the ones that only require the runtime API for this very purpose. 
created_at: 2025-05-13T20:22:53Z
author_association: MEMBER
diff_hunk: @@ -139,6 +139,36 @@ void _extend(cuvsResources_t res,
 
   cuvs::neighbors::ivf_flat::extend(*res_ptr, vectors_mds, indices_mds, index_ptr);
 }
+
+template <typename output_mdspan_type, typename T, typename IdxT>
+void _get_centers(cuvsResources_t res, cuvsIvfFlatIndex index, DLManagedTensor* centers)
+{
+  auto res_ptr   = reinterpret_cast<raft::resources*>(res);
+  auto index_ptr = reinterpret_cast<cuvs::neighbors::ivf_flat::index<T, IdxT>*>(index.addr);
+  auto dst       = cuvs::core::from_dlpack<output_mdspan_type>(centers);
+  auto src       = index_ptr->centers();
+
+  RAFT_EXPECTS(src.extent(0) == dst.extent(0), "Output centers has incorrect number of rows");
+  RAFT_EXPECTS(src.extent(1) == dst.extent(1), "Output centers has incorrect number of cols");
+
+  cudaMemcpyAsync(dst.data_handle(),
path: cpp/src/neighbors/ivf_flat_c.cpp
user: login: benfred
body: I updated nn_descent_c.cpp in the last commit to use raft::copy.

the `cudaMemcpyAsync` code is fine (as is copy functions from `raft/util/cudart_utils.hpp`) - I just couldn't use this code https://github.com/rapidsai/raft/blob/c2dc3124ce3fbcb5ff2ccabd88d7f57570b6aea9/cpp/include/raft/core/copy.cuh#L57-L61  from raft . 

fwiw, one nice thing about using raft::copy is that it fixes one issue that this code used to have (wasn't checking the return value from cudaMemCpyAsync , which was a stupid oversight on my part =) ).
created_at: 2025-05-13T20:27:01Z
author_association: CONTRIBUTOR
diff_hunk: @@ -311,6 +311,27 @@ void extract_centers(raft::resources const& res,
                                   stream));
 }
 
+void extract_centers(raft::resources const& res,
path: cpp/src/neighbors/ivf_pq/ivf_pq_build_common.cu
user: login: tarang-jain
body: Since we need both -- the host and device mdspan functions for `extract_centers`, we can have a pointer based function in the detail namespace, rather than duplicating the function for host mdspan. Then the host and device mdspan functions in the src file can call that pointer based function.
created_at: 2025-05-13T20:44:21Z
author_association: CONTRIBUTOR
diff_hunk: @@ -311,6 +311,27 @@ void extract_centers(raft::resources const& res,
                                   stream));
 }
 
+void extract_centers(raft::resources const& res,
path: cpp/src/neighbors/ivf_pq/ivf_pq_build_common.cu
user: login: cjnolet
body: I don't think we are going to be able to avoid this distinction here, and it's actually very low impact here to have the branched logic (in general, we want to move away from pointer-based functions across the codebase, not just in the public APIs... but we started w/ the public APIs). Not all device functions are super expensive.
created_at: 2025-05-13T20:47:47Z
author_association: MEMBER
diff_hunk: @@ -311,6 +311,27 @@ void extract_centers(raft::resources const& res,
                                   stream));
 }
 
+void extract_centers(raft::resources const& res,
path: cpp/src/neighbors/ivf_pq/ivf_pq_build_common.cu
user: login: tarang-jain
body: I am concerned that the content of the function is an exact copy except the function signature. In some other places we have had public APIs that are agnostic of what mdspan is being used (so the function argument is`raft::mdspan<const T, raft::matrix_extent<int64_t>, raft::row_major, data_accessor>`)
created_at: 2025-05-13T20:56:32Z
author_association: CONTRIBUTOR
diff_hunk: @@ -311,6 +311,27 @@ void extract_centers(raft::resources const& res,
                                   stream));
 }
 
+void extract_centers(raft::resources const& res,
path: cpp/src/neighbors/ivf_pq/ivf_pq_build_common.cu
user: login: benfred
body: I changed in the last commit to reduce the cut and paste code here (not using a pointer based api - instead just using a template version on the output mdspan type).  lmk what you all think
created_at: 2025-05-13T20:56:59Z
author_association: CONTRIBUTOR
diff_hunk: @@ -311,6 +311,27 @@ void extract_centers(raft::resources const& res,
                                   stream));
 }
 
+void extract_centers(raft::resources const& res,
path: cpp/src/neighbors/ivf_pq/ivf_pq_build_common.cu
user: login: tarang-jain
body: On the other hand, if we are okay with the two functions because it is not too expensive, that is also fine :))
created_at: 2025-05-13T20:57:43Z
author_association: CONTRIBUTOR
diff_hunk: @@ -311,6 +311,27 @@ void extract_centers(raft::resources const& res,
                                   stream));
 }
 
+void extract_centers(raft::resources const& res,
path: cpp/src/neighbors/ivf_pq/ivf_pq_build_common.cu
user: login: cjnolet
body: > I am concerned that the content of the function is an exact copy except the function signature. In some other places we have had public APIs that are agnostic of what mdspan is being used (so the function argument israft::mdspan<const T, raft::matrix_extent<int64_t>, raft::row_major, data_accessor>)

I'm in favor of accepting general `mdspan`. Not in favor of accepting raw pointers. Raw pointers are prone to too much error- it would be a giant step backwards for our code. 
created_at: 2025-05-13T20:58:01Z
author_association: MEMBER
diff_hunk: @@ -311,6 +311,27 @@ void extract_centers(raft::resources const& res,
                                   stream));
 }
 
+void extract_centers(raft::resources const& res,
path: cpp/src/neighbors/ivf_pq/ivf_pq_build_common.cu
user: login: benfred
body: this commit here is the one I meant https://github.com/rapidsai/cuvs/pull/888/commits/7080e2628b61d8c6fe6e2e4621b00ef2cad4f292 (takes a general mdspan )
created_at: 2025-05-13T20:59:11Z
author_association: CONTRIBUTOR
diff_hunk: @@ -311,6 +311,27 @@ void extract_centers(raft::resources const& res,
                                   stream));
 }
 
+void extract_centers(raft::resources const& res,
path: cpp/src/neighbors/ivf_pq/ivf_pq_build_common.cu
user: login: tarang-jain
body: Looks good. Just one nit pick (feel free to ignore since it is not that important). We can move detail namespace stuff to `ivf_pq_build.cuh`
created_at: 2025-05-13T21:08:00Z
author_association: CONTRIBUTOR
diff_hunk: @@ -311,6 +311,27 @@ void extract_centers(raft::resources const& res,
                                   stream));
 }
 
+void extract_centers(raft::resources const& res,
path: cpp/src/neighbors/ivf_pq/ivf_pq_build_common.cu
user: login: benfred
body: > Looks good. Just one nit pick (feel free to ignore since it is not that important). We can move detail namespace stuff to `ivf_pq_build.cuh`

moved in ed474d1
created_at: 2025-05-13T22:34:15Z
author_association: CONTRIBUTOR
diff_hunk: @@ -264,15 +264,16 @@ cuvsError_t cuvsIvfPqIndexDestroy(cuvsIvfPqIndex_t index);
 /** Get the number of clusters/inverted lists */
 uint32_t cuvsIvfPqIndexGetNLists(cuvsIvfPqIndex_t index);
 
-/** Get the dimensionality of the cluster centers */
-uint32_t cuvsIvfPqIndexGetDimExt(cuvsIvfPqIndex_t index);
+/** Get the dimensionality*/
path: cpp/include/cuvs/neighbors/ivf_pq.h
user: login: mythrocks
body: ```suggestion
/** Get the dimensionality */
```
created_at: 2025-05-19T17:09:39Z
author_association: CONTRIBUTOR
diff_hunk: @@ -351,3 +381,68 @@ extern "C" cuvsError_t cuvsIvfFlatExtend(cuvsResources_t res,
     }
   });
 }
+
+extern "C" uint32_t cuvsIvfFlatIndexGetNLists(cuvsIvfFlatIndex_t index)
+{
+  if (index->dtype.code == kDLFloat && index->dtype.bits == 32) {
+    auto index_ptr =
+      reinterpret_cast<cuvs::neighbors::ivf_flat::index<float, int64_t>*>(index->addr);
+    return index_ptr->n_lists();
+  } else if (index->dtype.code == kDLFloat && index->dtype.bits == 16) {
+    auto index_ptr =
+      reinterpret_cast<cuvs::neighbors::ivf_flat::index<half, int64_t>*>(index->addr);
+    return index_ptr->n_lists();
+  } else if (index->dtype.code == kDLInt && index->dtype.bits == 8) {
+    auto index_ptr =
+      reinterpret_cast<cuvs::neighbors::ivf_flat::index<int8_t, int64_t>*>(index->addr);
+    return index_ptr->n_lists();
+  } else if (index->dtype.code == kDLUInt && index->dtype.bits == 8) {
+    auto index_ptr =
+      reinterpret_cast<cuvs::neighbors::ivf_flat::index<uint8_t, int64_t>*>(index->addr);
+    return index_ptr->n_lists();
path: cpp/src/neighbors/ivf_flat_c.cpp
user: login: mythrocks
body: This is a recurring pattern in the code.  One wonders if `libcudf`'s type-dispatch pattern might be of value here.

We might consider exploring at a later date.
created_at: 2025-05-19T17:54:57Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1904,4 +1904,26 @@ void extend(
                   new_indices.has_value() ? new_indices.value().data_handle() : nullptr,
                   n_rows);
 }
+
+template <typename output_mdspan_type>
+inline void extract_centers(raft::resources const& res,
+                            const cuvs::neighbors::ivf_pq::index<int64_t>& index,
+                            output_mdspan_type cluster_centers)
+{
+  RAFT_EXPECTS(cluster_centers.extent(0) == index.n_lists(),
+               "Number of rows in the output buffer for cluster centers must be equal to the "
+               "number of IVF lists");
+  RAFT_EXPECTS(
+    cluster_centers.extent(1) == index.dim(),
+    "Number of columns in the output buffer for cluster centers and index dim are different");
+  auto stream = raft::resource::get_cuda_stream(res);
+  RAFT_CUDA_TRY(cudaMemcpy2DAsync(cluster_centers.data_handle(),
+                                  sizeof(float) * index.dim(),
+                                  index.centers().data_handle(),
+                                  sizeof(float) * index.dim_ext(),
+                                  sizeof(float) * index.dim(),
+                                  index.n_lists(),
+                                  cudaMemcpyDefault,
path: cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh
user: login: mythrocks
body: TIL `cudaMemcpyDefault`.  I didn't know/realize that the direction could be inferred.
created_at: 2025-05-19T18:14:17Z
author_association: CONTRIBUTOR
diff_hunk: @@ -139,6 +140,35 @@ void _extend(cuvsResources_t res,
 
   cuvs::neighbors::ivf_flat::extend(*res_ptr, vectors_mds, indices_mds, index_ptr);
 }
+
+template <typename output_mdspan_type, typename T, typename IdxT>
+void _get_centers(cuvsResources_t res, cuvsIvfFlatIndex index, DLManagedTensor* centers)
+{
+  auto res_ptr   = reinterpret_cast<raft::resources*>(res);
+  auto index_ptr = reinterpret_cast<cuvs::neighbors::ivf_flat::index<T, IdxT>*>(index.addr);
+  auto dst       = cuvs::core::from_dlpack<output_mdspan_type>(centers);
+  auto src       = index_ptr->centers();
path: cpp/src/neighbors/ivf_flat_c.cpp
user: login: mythrocks
body: Would there be value in making any of these `const`?  

I don't have familiarity with the code yet, so I haven't grokked the semantics of `raft::copy`.  Apologies, if this is noise.
created_at: 2025-05-19T18:33:56Z
author_association: CONTRIBUTOR

PR comments:
user: login: cjnolet
created_at: 2025-05-27T23:01:43Z
body: /merge
author_association: MEMBER