IVF-Flat reconstruction

cpp/CMakeLists.txt

@@ -337,6 +337,9 @@ if(RAFT_COMPILE_LIBRARY)
     src/neighbors/ivf_flat_search_float_int64_t.cu
     src/neighbors/ivf_flat_search_int8_t_int64_t.cu
     src/neighbors/ivf_flat_search_uint8_t_int64_t.cu
+    src/neighbors/ivf_flat_reconstruct_float_int64_t.cu
+    src/neighbors/ivf_flat_reconstruct_int8_t_int64_t.cu
+    src/neighbors/ivf_flat_reconstruct_uint8_t_int64_t.cu
     src/neighbors/ivfpq_build_float_int64_t.cu
     src/neighbors/ivfpq_build_int8_t_int64_t.cu
     src/neighbors/ivfpq_build_uint8_t_int64_t.cu
@@ -602,7 +605,9 @@ target_link_libraries(raft::raft INTERFACE
 # Use `rapids_export` for 22.04 as it will have COMPONENT support
 rapids_export(
   INSTALL raft
-  EXPORT_SET raft-exports COMPONENTS ${raft_components} COMPONENTS_EXPORT_SET ${raft_export_sets}
+  EXPORT_SET raft-exports
+  COMPONENTS ${raft_components}
+  COMPONENTS_EXPORT_SET ${raft_export_sets}
   GLOBAL_TARGETS raft compiled distributed
   NAMESPACE raft::
   DOCUMENTATION doc_string
@@ -613,7 +618,9 @@ rapids_export(
 # * build export -------------------------------------------------------------
 rapids_export(
   BUILD raft
-  EXPORT_SET raft-exports COMPONENTS ${raft_components} COMPONENTS_EXPORT_SET ${raft_export_sets}
+  EXPORT_SET raft-exports
+  COMPONENTS ${raft_components}
+  COMPONENTS_EXPORT_SET ${raft_export_sets}
   GLOBAL_TARGETS raft compiled distributed
   DOCUMENTATION doc_string
   NAMESPACE raft::

cpp/bench/ann/src/faiss/faiss_benchmark.cu

@@ -104,10 +104,10 @@ std::unique_ptr<raft::bench::ann::ANN<T>> create_algo(const std::string& algo,
   // stop compiler warning; not all algorithms support multi-GPU so it may not be used
   (void)dev_list;
 
-  raft::bench::ann::Metric metric = parse_metric(distance);
   std::unique_ptr<raft::bench::ann::ANN<T>> ann;
 
   if constexpr (std::is_same_v<T, float>) {
+    raft::bench::ann::Metric metric = parse_metric(distance);
     if (algo == "faiss_gpu_ivf_flat") {
       ann = make_algo<T, raft::bench::ann::FaissGpuIVFFlat>(metric, dim, conf, dev_list);
     } else if (algo == "faiss_gpu_ivf_pq") {

cpp/include/raft/neighbors/detail/ivf_flat_build.cuh

@@ -17,11 +17,13 @@
 #pragma once
 
 #include <raft/cluster/kmeans_balanced.cuh>
+#include <raft/core/device_mdspan.hpp>
 #include <raft/core/logger.hpp>
 #include <raft/core/mdarray.hpp>
 #include <raft/core/nvtx.hpp>
 #include <raft/core/operators.hpp>
 #include <raft/core/resource/cuda_stream.hpp>
+#include <raft/core/resource/thrust_policy.hpp>
 #include <raft/core/resources.hpp>
 #include <raft/linalg/add.cuh>
 #include <raft/linalg/map.cuh>
@@ -35,6 +37,8 @@
 
 #include <rmm/cuda_stream_view.hpp>
 
+#include <thrust/extrema.h>
+
 #include <cstdint>
 
 namespace raft::neighbors::ivf_flat::detail {
@@ -416,4 +420,152 @@ inline void fill_refinement_index(raft::resources const& handle,
                                          refinement_index->veclen());
   RAFT_CUDA_TRY(cudaPeekAtLastError());
 }
+
+template <typename T, typename IdxT>
+__global__ void get_data_ptr_kernel(const uint32_t* list_sizes,
+                                    const T* const* list_data_ptrs,
+                                    const IdxT* const* list_indices_ptrs,
+                                    uint32_t dim,
+                                    uint32_t veclen,
+                                    uint32_t n_list,
+                                    IdxT max_indice,
+                                    T** ptrs_to_data)
+{
+  const IdxT list_id = IdxT(blockDim.x) * IdxT(blockIdx.x) + threadIdx.x;
+  if (list_id >= n_list) { return; }
+  const IdxT inlist_id     = IdxT(blockDim.y) * IdxT(blockIdx.y) + threadIdx.y;
+  const uint32_t list_size = list_sizes[list_id];
+  if (inlist_id >= list_size) { return; }
+
+  auto* list_indices = list_indices_ptrs[list_id];
+  IdxT id            = list_indices[inlist_id];
+  if (id > max_indice) { return; }
+
+  using interleaved_group = Pow2<kIndexGroupSize>;
+  auto group_offset       = interleaved_group::roundDown(inlist_id);
+  auto ingroup_id         = interleaved_group::mod(inlist_id) * veclen;
+
+  auto* list_data  = list_data_ptrs[list_id];
+  const T* ptr     = list_data + (group_offset * dim) + ingroup_id;
+  ptrs_to_data[id] = (T*)ptr;
+}
+
+template <typename T, typename IdxT>
+__global__ void reconstruct_batch_kernel(const IdxT* vector_ids,
+                                         const T** ptrs_to_data,
+                                         uint32_t dim,
+                                         uint32_t veclen,
+                                         IdxT n_rows,
+                                         T* reconstr)
+{
+  const IdxT i = IdxT(blockDim.x) * IdxT(blockIdx.x) + threadIdx.x;
+  if (i >= n_rows) { return; }
+
+  const IdxT vector_id = vector_ids[i];
+  const T* src         = ptrs_to_data[vector_id];
+  if (!src) { return; }
+
+  reconstr += i * dim;
+  for (uint32_t l = 0; l < dim; l += veclen) {
+    for (uint32_t j = 0; j < veclen; j++) {
+      reconstr[l + j] = src[l * kIndexGroupSize + j];
+    }
+  }
+}
+
+template <typename T, typename IdxT>
+void reconstruct_batch(raft::resources const& handle,
+                       const index<T, IdxT>& index,
+                       raft::device_vector_view<const IdxT, IdxT> vector_ids,
+                       raft::device_matrix_view<T, IdxT, row_major> vector_out)
+{
+  auto stream      = raft::resource::get_cuda_stream(handle);
+  auto exec_policy = raft::resource::get_thrust_policy(handle);
+
+  thrust::device_ptr<const IdxT> vector_ids_ptr =
+    thrust::device_pointer_cast(vector_ids.data_handle());
+  IdxT max_indice =
+    *thrust::max_element(exec_policy, vector_ids_ptr, vector_ids_ptr + vector_ids.extent(0));
+
+  rmm::device_uvector<T*> ptrs_to_data(max_indice + 1, stream);
+  utils::memzero(ptrs_to_data.data(), ptrs_to_data.size(), stream);
+
+  thrust::device_ptr<const uint32_t> list_sizes_ptr =
+    thrust::device_pointer_cast(index.list_sizes().data_handle());
+  uint32_t max_list_size = *thrust::max_element(
+    exec_policy, list_sizes_ptr, list_sizes_ptr + index.list_sizes().extent(0));
+
+  const dim3 block_dim1(16, 16);
+  const dim3 grid_dim1(raft::ceildiv<size_t>(index.n_lists(), block_dim1.x),
+                       raft::ceildiv<size_t>(max_list_size, block_dim1.y));
+  get_data_ptr_kernel<<<grid_dim1, block_dim1, 0, stream>>>(index.list_sizes().data_handle(),
+                                                            index.data_ptrs().data_handle(),
+                                                            index.inds_ptrs().data_handle(),
+                                                            index.dim(),
+                                                            index.veclen(),
+                                                            index.n_lists(),
+                                                            max_indice,
+                                                            ptrs_to_data.data());
+  RAFT_CUDA_TRY(cudaPeekAtLastError());
+
+  auto n_reconstruction = vector_ids.extent(0);
+  const dim3 block_dim2(256);
+  const dim3 grid_dim2(raft::ceildiv<size_t>(n_reconstruction, block_dim2.x));
+  reconstruct_batch_kernel<<<grid_dim2, block_dim2, 0, stream>>>(vector_ids.data_handle(),
+                                                                 (const T**)ptrs_to_data.data(),
+                                                                 index.dim(),
+                                                                 index.veclen(),
+                                                                 n_reconstruction,
+                                                                 vector_out.data_handle());
+  RAFT_CUDA_TRY(cudaPeekAtLastError());
+}
+
+template <typename T>
+__global__ void reconstruct_list_data_kernel(
+  T* out_vectors,
+  T* in_list_data,
+  std::variant<uint32_t, const uint32_t*> offset_or_indices,
+  uint32_t len,
+  size_t veclen,
+  uint32_t dim)
+{
+  for (uint32_t ix = threadIdx.x + blockDim.x * blockIdx.x; ix < len; ix += blockDim.x) {
+    const uint32_t src_ix = std::holds_alternative<uint32_t>(offset_or_indices)
+                              ? std::get<uint32_t>(offset_or_indices) + ix
+                              : std::get<const uint32_t*>(offset_or_indices)[ix];
+
+    using group_align         = Pow2<kIndexGroupSize>;
+    const uint32_t group_ix   = group_align::div(src_ix);
+    const uint32_t ingroup_ix = group_align::mod(src_ix) * veclen;
+
+    for (uint32_t l = 0; l < dim; l += veclen) {
+      for (uint32_t j = 0; j < veclen; j++) {
+        out_vectors[ix * dim + l + j] = in_list_data[l * kIndexGroupSize + ingroup_ix + j];
+      }
+    }
+  }
+}
+
+/** Decode the list data; see the public interface for the api and usage. */
+template <typename T, typename IdxT>
+void reconstruct_list_data(raft::resources const& handle,
+                           const index<T, IdxT>& index,
+                           device_matrix_view<T, IdxT, row_major> out_vectors,
+                           uint32_t label,
+                           uint32_t offset)
+{
+  auto stream = raft::resource::get_cuda_stream(handle);
+
+  uint32_t len = out_vectors.extent(0);
+  const dim3 block_dim(256);
+  const dim3 grid_dim(raft::ceildiv<size_t>(len, block_dim.x));
+  reconstruct_list_data_kernel<<<grid_dim, block_dim, 0, stream>>>(
+    out_vectors.data_handle(),
+    index.lists()[label]->data.data_handle(),
+    offset,
+    len,
+    (size_t)index.veclen(),
+    index.dim());
+}
+
 }  // namespace raft::neighbors::ivf_flat::detail

cpp/include/raft/neighbors/ivf_flat-ext.cuh

@@ -93,6 +93,19 @@ void search(raft::resources const& handle,
             raft::device_matrix_view<IdxT, IdxT, row_major> neighbors,
             raft::device_matrix_view<float, IdxT, row_major> distances) RAFT_EXPLICIT;
 
+template <typename T, typename IdxT>
+void reconstruct_batch(raft::resources const& handle,
+                       const index<T, IdxT>& index,
+                       raft::device_vector_view<const IdxT, IdxT> vector_ids,
+                       raft::device_matrix_view<T, IdxT, row_major> vector_out) RAFT_EXPLICIT;
+
+template <typename T, typename IdxT>
+void reconstruct_list_data(raft::resources const& handle,
+                           const index<T, IdxT>& index,
+                           device_matrix_view<T, IdxT, row_major> out_vectors,
+                           uint32_t label,
+                           uint32_t offset) RAFT_EXPLICIT;
+
 }  // namespace raft::neighbors::ivf_flat
 
 #endif  // RAFT_EXPLICIT_INSTANTIATE_ONLY
@@ -183,3 +196,23 @@ instantiate_raft_neighbors_ivf_flat_search(int8_t, int64_t);
 instantiate_raft_neighbors_ivf_flat_search(uint8_t, int64_t);
 
 #undef instantiate_raft_neighbors_ivf_flat_search
+
+#define instantiate_raft_neighbors_ivf_flat_reconstruct(T, IdxT)                  \
+  extern template void raft::neighbors::ivf_flat::reconstruct_batch<T, IdxT>(     \
+    raft::resources const& handle,                                                \
+    const raft::neighbors::ivf_flat::index<T, IdxT>& index,                       \
+    raft::device_vector_view<const IdxT, IdxT> vector_ids,                        \
+    raft::device_matrix_view<T, IdxT, row_major> vector_out);                     \
+                                                                                  \
+  extern template void raft::neighbors::ivf_flat::reconstruct_list_data<T, IdxT>( \
+    raft::resources const& handle,                                                \
+    const raft::neighbors::ivf_flat::index<T, IdxT>& index,                       \
+    raft::device_matrix_view<T, IdxT, row_major> out_vectors,                     \
+    uint32_t label,                                                               \
+    uint32_t offset);
+
+instantiate_raft_neighbors_ivf_flat_reconstruct(float, int64_t);
+instantiate_raft_neighbors_ivf_flat_reconstruct(int8_t, int64_t);
+instantiate_raft_neighbors_ivf_flat_reconstruct(uint8_t, int64_t);
+
+#undef instantiate_raft_neighbors_ivf_flat_reconstruct

cpp/include/raft/neighbors/ivf_flat-inl.cuh

@@ -464,6 +464,52 @@ void search(raft::resources const& handle,
                 nullptr);
 }
 
+/**
+ * @brief Reconstruct vector database
+ *
+ *
+ * @tparam T data element type
+ * @tparam IdxT type of the indices
+ *
+ * @param[in] handle
+ * @param[in] index ivf-flat constructed index
+ * @param[in] vector_ids vector with the ids of vectors to look for
+ * @param[out] vector_out matrix with the database of vectors
+ */
+template <typename T, typename IdxT>
+void reconstruct_batch(raft::resources const& handle,
+                       const index<T, IdxT>& index,
+                       raft::device_vector_view<const IdxT, IdxT> vector_ids,
+                       raft::device_matrix_view<T, IdxT, row_major> vector_out)
+{
+  return raft::neighbors::ivf_flat::detail::reconstruct_batch(
+    handle, index, vector_ids, vector_out);
+}
+
+/**
+ * @brief Reconstruct vectors of a given cluster
+ *
+ *
+ * @tparam T data element type
+ * @tparam IdxT type of the indices
+ *
+ * @param[in] handle
+ * @param[in] index ivf-flat constructed index
+ * @param[out] vector_out matrix with the vectors contained in the cluster
+ * @param[in] label cluster index
+ * @param[in] offset offset
+ */
+template <typename T, typename IdxT>
+void reconstruct_list_data(raft::resources const& handle,
+                           const index<T, IdxT>& index,
+                           device_matrix_view<T, IdxT, row_major> out_vectors,
+                           uint32_t label,
+                           uint32_t offset)
+{
+  return raft::neighbors::ivf_flat::detail::reconstruct_list_data(
+    handle, index, out_vectors, label, offset);
+}
+
 /** @} */
 
 }  // namespace raft::neighbors::ivf_flat

cpp/include/raft_runtime/neighbors/ivf_flat.hpp

@@ -46,6 +46,17 @@ namespace raft::runtime::neighbors::ivf_flat {
               std::optional<raft::device_vector_view<const IdxT, IdxT>> new_indices, \
               raft::neighbors::ivf_flat::index<T, IdxT>* idx);                       \
                                                                                      \
+  void reconstruct_batch(raft::resources const& handle,                              \
+                         raft::neighbors::ivf_flat::index<T, IdxT>& idx,             \
+                         raft::device_vector_view<const IdxT, IdxT> vector_ids,      \
+                         raft::device_matrix_view<T, IdxT, row_major> vector_out);   \
+                                                                                     \
+  void reconstruct_list_data(raft::resources const& handle,                          \
+                             const raft::neighbors::ivf_flat::index<T, IdxT>& idx,   \
+                             device_matrix_view<T, IdxT, row_major> out_vectors,     \
+                             uint32_t label,                                         \
+                             uint32_t offset);                                       \
+                                                                                     \
   void serialize_file(raft::resources const& handle,                                 \
                       const std::string& filename,                                   \
                       const raft::neighbors::ivf_flat::index<T, IdxT>& index);       \

cpp/src/neighbors/ivf_flat_reconstruct_float_int64_t.cu

@@ -0,0 +1,44 @@
+/*
+ * Copyright (c) 2023, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+/*
+ * NOTE: this file is generated by ivf_flat_00_generate.py
+ *
+ * Make changes there and run in this directory:
+ *
+ * > python ivf_flat_00_generate.py
+ *
+ */
+
+#include <raft/neighbors/ivf_flat-inl.cuh>
+
+#define instantiate_raft_neighbors_ivf_flat_reconstruct(T, IdxT)           \
+  template void raft::neighbors::ivf_flat::reconstruct_batch<T, IdxT>(     \
+    raft::resources const& handle,                                         \
+    const raft::neighbors::ivf_flat::index<T, IdxT>& idx,                  \
+    raft::device_vector_view<const IdxT, IdxT> vector_ids,                 \
+    raft::device_matrix_view<T, IdxT, row_major> vector_out);              \
+                                                                           \
+  template void raft::neighbors::ivf_flat::reconstruct_list_data<T, IdxT>( \
+    raft::resources const& handle,                                         \
+    const raft::neighbors::ivf_flat::index<T, IdxT>& idx,                  \
+    raft::device_matrix_view<T, IdxT, row_major> out_vectors,              \
+    uint32_t label,                                                        \
+    uint32_t offset);
+
+instantiate_raft_neighbors_ivf_flat_reconstruct(float, int64_t);
+
+#undef instantiate_raft_neighbors_ivf_flat_reconstruct