Merge branch 'c-add-memory-pool' of https://github.com/ajit283/cuvs into c-add-memory-pool

Author: ajit283

.gitignore

@@ -74,3 +74,4 @@ cagra_index
 ivf_flat_index
 ivf_pq_index
 
+.aider*

build.sh

@@ -18,14 +18,15 @@ ARGS=$*
 # scripts, and that this script resides in the repo dir!
 REPODIR=$(cd $(dirname $0); pwd)
 
-VALIDARGS="clean libcuvs python rust docs tests bench-ann examples --uninstall  -v -g -n --compile-static-lib --allgpuarch --no-nvtx --show_depr_warn --incl-cache-stats --time -h"
+VALIDARGS="clean libcuvs python rust go docs tests bench-ann examples --uninstall  -v -g -n --compile-static-lib --allgpuarch --no-nvtx --show_depr_warn --incl-cache-stats --time -h"
 HELP="$0 [<target> ...] [<flag> ...] [--cmake-args=\"<args>\"] [--cache-tool=<tool>] [--limit-tests=<targets>] [--limit-bench-ann=<targets>] [--build-metrics=<filename>]
  where <target> is:
    clean            - remove all existing build artifacts and configuration (start over)
    libcuvs          - build the cuvs C++ code only. Also builds the C-wrapper library
                       around the C++ code.
    python           - build the cuvs Python package
    rust             - build the cuvs Rust bindings
+   go               - build the cuvs Go bindings
    docs             - build the documentation
    tests            - build the tests
    bench-ann        - build end-to-end ann benchmarks
@@ -426,6 +427,13 @@ if (( ${NUMARGS} == 0 )) || hasArg rust; then
     cargo test
 fi
 
+# Build the cuvs Go bindings
+if (( ${NUMARGS} == 0 )) || hasArg go; then
+    cd ${REPODIR}/go
+    go build ./...
+    go test ./...
+fi
+
 export RAPIDS_VERSION="$(sed -E -e 's/^([0-9]{2})\.([0-9]{2})\.([0-9]{2}).*$/\1.\2.\3/' "${REPODIR}/VERSION")"
 export RAPIDS_VERSION_MAJOR_MINOR="$(sed -E -e 's/^([0-9]{2})\.([0-9]{2})\.([0-9]{2}).*$/\1.\2/' "${REPODIR}/VERSION")"
 

ci/build_go.sh

@@ -0,0 +1,33 @@
+#!/bin/bash
+# Copyright (c) 2024, NVIDIA CORPORATION.
+
+set -euo pipefail
+
+rapids-logger "Create test conda environment"
+. /opt/conda/etc/profile.d/conda.sh
+
+
+rapids-mamba-retry env create --yes -n go
+
+# seeing failures on activating the environment here on unbound locals
+# apply workaround from https://github.com/conda/conda/issues/8186#issuecomment-532874667
+set +eu
+conda activate go
+set -eu
+
+rapids-print-env
+
+export CGO_CFLAGS="-I$CONDA_PREFIX/include -I/usr/local/cuda/include  -I/usr/local/include"
+export CGO_LDFLAGS="-L$CONDA_PREFIX/lib -L/usr/local/cuda/lib64  -lcuvs_c -lcudart     -Wl,-rpath,$CONDA_PREFIX/lib-Wl,-rpath,/usr/local/cuda/lib64"
+
+rapids-logger "Downloading artifacts from previous jobs"
+CPP_CHANNEL=$(rapids-download-conda-from-s3 cpp)
+
+# installing libcuvs/libraft will speed up the rust build substantially
+rapids-mamba-retry install \
+  --channel "${CPP_CHANNEL}" \
+  libcuvs  \
+  libraft  \
+  cuvs
+
+bash ./build.sh go

cpp/include/cuvs/core/c_api.h

@@ -140,6 +140,7 @@ cuvsError_t cuvsRMMFree(cuvsResources_t res, void* ptr, size_t bytes);
  * available memory
  * @return cuvsError_t
  */
+
 cuvsError_t cuvsRMMPoolMemoryResourceEnable(int initial_pool_size_percent,
                                             int max_pool_size_percent,
                                             int managed);
@@ -149,6 +150,7 @@ cuvsError_t cuvsRMMPoolMemoryResourceEnable(int initial_pool_size_percent,
  */
 cuvsError_t cuvsRMMMemoryResourceReset();
 
+
 /** @} */
 
 #ifdef __cplusplus

cpp/include/cuvs/neighbors/cagra.h

@@ -137,6 +137,42 @@ cuvsError_t cuvsCagraCompressionParamsCreate(cuvsCagraCompressionParams_t* param
  */
 cuvsError_t cuvsCagraCompressionParamsDestroy(cuvsCagraCompressionParams_t params);
 
+/**
+ * @}
+ */
+
+/**
+ * @defgroup cagra_c_extend_params C API for CUDA ANN Graph-based nearest neighbor search
+ * @{
+ */
+/**
+ * @brief Supplemental parameters to extend CAGRA Index
+ *
+ */
+struct cuvsCagraExtendParams {
+  /** Degree of input graph for pruning. */
+  uint32_t max_chunk_size;
+  /** Degree of output graph. */
+};
+
+typedef struct cuvsCagraExtendParams* cuvsCagraExtendParams_t;
+
+/**
+ * @brief Allocate CAGRA Extend params, and populate with default values
+ *
+ * @param[in] params cuvsCagraExtendParams_t to allocate
+ * @return cuvsError_t
+ */
+cuvsError_t cuvsCagraExtendParamsCreate(cuvsCagraExtendParams_t* params);
+
+/**
+ * @brief De-allocate CAGRA Extend params
+ *
+ * @param[in] params
+ * @return cuvsError_t
+ */
+cuvsError_t cuvsCagraExtendParamsDestroy(cuvsCagraExtendParams_t params);
+
 /**
  * @}
  */
@@ -327,6 +363,34 @@ cuvsError_t cuvsCagraBuild(cuvsResources_t res,
  * @}
  */
 
+/**
+ * @defgroup cagra_c_extend_params C API for CUDA ANN Graph-based nearest neighbor search
+ * @{
+ */
+
+/**
+ * @brief Extend a CAGRA index with a `DLManagedTensor` which has underlying
+ *        `DLDeviceType` equal to `kDLCUDA`, `kDLCUDAHost`, `kDLCUDAManaged`,
+ *        or `kDLCPU`. Also, acceptable underlying types are:
+ *        1. `kDLDataType.code == kDLFloat` and `kDLDataType.bits = 32`
+ *        2. `kDLDataType.code == kDLInt` and `kDLDataType.bits = 8`
+ *        3. `kDLDataType.code == kDLUInt` and `kDLDataType.bits = 8`
+ *
+ * @param[in] res cuvsResources_t opaque C handle
+ * @param[in] params cuvsCagraExtendParams_t used to extend CAGRA index
+ * @param[in] additional_dataset DLManagedTensor* additional dataset
+ * @param[out] index cuvsCagraIndex_t Extended CAGRA index
+ * @return cuvsError_t
+ */
+cuvsError_t cuvsCagraExtend(cuvsResources_t res,
+                            cuvsCagraExtendParams_t params,
+                            DLManagedTensor* additional_dataset,
+                            cuvsCagraIndex_t index);
+
+/**
+ * @}
+ */
+
 /**
  * @defgroup cagra_c_index_search C API for CUDA ANN Graph-based nearest neighbor search
  * @{

cpp/src/core/c_api.cpp

@@ -26,7 +26,6 @@
 #include <rmm/mr/device/owning_wrapper.hpp>
 #include <rmm/mr/device/per_device_resource.hpp>
 #include <rmm/mr/device/pool_memory_resource.hpp>
-
 #include <thread>
 
 extern "C" cuvsError_t cuvsResourcesCreate(cuvsResources_t* res)
@@ -69,6 +68,8 @@ extern "C" cuvsError_t cuvsStreamSync(cuvsResources_t res)
   });
 }
 
+thread_local std::unique_ptr<rmm::mr::pool_memory_resource<rmm::mr::cuda_memory_resource>> pool_mr;
+
 extern "C" cuvsError_t cuvsRMMAlloc(cuvsResources_t res, void** ptr, size_t bytes)
 {
   return cuvs::core::translate_exceptions([=] {
@@ -87,6 +88,7 @@ extern "C" cuvsError_t cuvsRMMFree(cuvsResources_t res, void* ptr, size_t bytes)
   });
 }
 
+
 thread_local std::shared_ptr<
   rmm::mr::owning_wrapper<rmm::mr::pool_memory_resource<rmm::mr::device_memory_resource>,
                           rmm::mr::device_memory_resource>>
@@ -104,6 +106,7 @@ extern "C" cuvsError_t cuvsRMMPoolMemoryResourceEnable(int initial_pool_size_per
       throw std::runtime_error("Current memory resource is not a cuda_memory_resource");
     }
 
+
     auto initial_size = rmm::percent_of_free_device_memory(initial_pool_size_percent);
     auto max_size     = rmm::percent_of_free_device_memory(max_pool_size_percent);
 
@@ -119,10 +122,12 @@ extern "C" cuvsError_t cuvsRMMPoolMemoryResourceEnable(int initial_pool_size_per
     pool_mr =
       rmm::mr::make_owning_wrapper<rmm::mr::pool_memory_resource>(mr, initial_size, max_size);
 
+
     rmm::mr::set_current_device_resource(pool_mr.get());
   });
 }
 
+
 extern "C" cuvsError_t cuvsRMMMemoryResourceReset()
 {
   return cuvs::core::translate_exceptions([=] {

cpp/src/neighbors/cagra_c.cpp

@@ -83,6 +83,34 @@ void* _build(cuvsResources_t res, cuvsCagraIndexParams params, DLManagedTensor*
   return index;
 }
 
+template <typename T>
+void _extend(cuvsResources_t res,
+             cuvsCagraExtendParams params,
+             cuvsCagraIndex index,
+             DLManagedTensor* additional_dataset_tensor)
+{
+  auto dataset   = additional_dataset_tensor->dl_tensor;
+  auto index_ptr = reinterpret_cast<cuvs::neighbors::cagra::index<T, uint32_t>*>(index.addr);
+  auto res_ptr   = reinterpret_cast<raft::resources*>(res);
+
+  auto extend_params           = cuvs::neighbors::cagra::extend_params();
+  extend_params.max_chunk_size = params.max_chunk_size;
+
+  if (cuvs::core::is_dlpack_device_compatible(dataset)) {
+    using mdspan_type = raft::device_matrix_view<T const, int64_t, raft::row_major>;
+    auto mds          = cuvs::core::from_dlpack<mdspan_type>(additional_dataset_tensor);
+    cuvs::neighbors::cagra::extend(*res_ptr, extend_params, mds, *index_ptr);
+  } else if (cuvs::core::is_dlpack_host_compatible(dataset)) {
+    using mdspan_type = raft::host_matrix_view<T const, int64_t, raft::row_major>;
+    auto mds          = cuvs::core::from_dlpack<mdspan_type>(additional_dataset_tensor);
+    cuvs::neighbors::cagra::extend(*res_ptr, extend_params, mds, *index_ptr);
+  } else {
+    RAFT_FAIL("Unsupported dataset DLtensor dtype: %d and bits: %d",
+              dataset.dtype.code,
+              dataset.dtype.bits);
+  }
+}
+
 template <typename T>
 void _search(cuvsResources_t res,
              cuvsCagraSearchParams params,
@@ -190,6 +218,29 @@ extern "C" cuvsError_t cuvsCagraBuild(cuvsResources_t res,
   });
 }
 
+extern "C" cuvsError_t cuvsCagraExtend(cuvsResources_t res,
+                                       cuvsCagraExtendParams_t params,
+                                       DLManagedTensor* additional_dataset_tensor,
+                                       cuvsCagraIndex_t index_c_ptr)
+{
+  return cuvs::core::translate_exceptions([=] {
+    auto dataset = additional_dataset_tensor->dl_tensor;
+    auto index   = *index_c_ptr;
+
+    if ((dataset.dtype.code == kDLFloat) && (dataset.dtype.bits == 32)) {
+      _extend<float>(res, *params, index, additional_dataset_tensor);
+    } else if (dataset.dtype.code == kDLInt && dataset.dtype.bits == 8) {
+      _extend<int8_t>(res, *params, index, additional_dataset_tensor);
+    } else if (dataset.dtype.code == kDLUInt && dataset.dtype.bits == 8) {
+      _extend<uint8_t>(res, *params, index, additional_dataset_tensor);
+    } else {
+      RAFT_FAIL("Unsupported dataset DLtensor dtype: %d and bits: %d",
+                dataset.dtype.code,
+                dataset.dtype.bits);
+    }
+  });
+}
+
 extern "C" cuvsError_t cuvsCagraSearch(cuvsResources_t res,
                                        cuvsCagraSearchParams_t params,
                                        cuvsCagraIndex_t index_c_ptr,
@@ -265,6 +316,17 @@ extern "C" cuvsError_t cuvsCagraCompressionParamsDestroy(cuvsCagraCompressionPar
   return cuvs::core::translate_exceptions([=] { delete params; });
 }
 
+extern "C" cuvsError_t cuvsCagraExtendParamsCreate(cuvsCagraExtendParams_t* params)
+{
+  return cuvs::core::translate_exceptions(
+    [=] { *params = new cuvsCagraExtendParams{.max_chunk_size = 0}; });
+}
+
+extern "C" cuvsError_t cuvsCagraExtendParamsDestroy(cuvsCagraExtendParams_t params)
+{
+  return cuvs::core::translate_exceptions([=] { delete params; });
+}
+
 extern "C" cuvsError_t cuvsCagraSearchParamsCreate(cuvsCagraSearchParams_t* params)
 {
   return cuvs::core::translate_exceptions([=] {

cpp/test/core/c_api.c

@@ -55,6 +55,7 @@ int main()
   if (free_error_pool == CUVS_ERROR) { exit(EXIT_FAILURE); }
 
   // Reset pool memory resource
+
   cuvsError_t reset_error = cuvsRMMMemoryResourceReset();
   if (reset_error == CUVS_ERROR) { exit(EXIT_FAILURE); }
 

cpp/test/neighbors/ann_cagra_c.cu

@@ -57,6 +57,20 @@ TEST(CagraC, BuildSearch)
   dataset_tensor.dl_tensor.shape              = dataset_shape;
   dataset_tensor.dl_tensor.strides            = nullptr;
 
+  // create additional dataset DLTensor
+  rmm::device_uvector<float> additional_d(4 * 2, stream);
+  raft::copy(additional_d.data(), (float*)dataset, 4 * 2, stream);
+  DLManagedTensor additional_dataset_tensor;
+  additional_dataset_tensor.dl_tensor.data               = additional_d.data();
+  additional_dataset_tensor.dl_tensor.device.device_type = kDLCUDA;
+  additional_dataset_tensor.dl_tensor.ndim               = 2;
+  additional_dataset_tensor.dl_tensor.dtype.code         = kDLFloat;
+  additional_dataset_tensor.dl_tensor.dtype.bits         = 32;
+  additional_dataset_tensor.dl_tensor.dtype.lanes        = 1;
+  int64_t additional_dataset_shape[2]                    = {4, 2};
+  additional_dataset_tensor.dl_tensor.shape              = additional_dataset_shape;
+  additional_dataset_tensor.dl_tensor.strides            = nullptr;
+
   // create index
   cuvsCagraIndex_t index;
   cuvsCagraIndexCreate(&index);
@@ -66,6 +80,11 @@ TEST(CagraC, BuildSearch)
   cuvsCagraIndexParamsCreate(&build_params);
   cuvsCagraBuild(res, build_params, &dataset_tensor, index);
 
+  cuvsCagraExtendParams_t extend_params;
+  cuvsCagraExtendParamsCreate(&extend_params);
+  extend_params->max_chunk_size = 100;
+  cuvsCagraExtend(res, extend_params, &additional_dataset_tensor, index);
+
   // create queries DLTensor
   rmm::device_uvector<float> queries_d(4 * 2, stream);
   raft::copy(queries_d.data(), (float*)queries, 4 * 2, stream);
@@ -122,6 +141,7 @@ TEST(CagraC, BuildSearch)
 
   // de-allocate index and res
   cuvsCagraSearchParamsDestroy(search_params);
+  cuvsCagraExtendParamsDestroy(extend_params);
   cuvsCagraIndexParamsDestroy(build_params);
   cuvsCagraIndexDestroy(index);
   cuvsResourcesDestroy(res);

examples/c/CMakeLists.txt

@@ -13,6 +13,8 @@
 
 cmake_minimum_required(VERSION 3.26.4 FATAL_ERROR)
 
+set(CMAKE_CUDA_ARCHITECTURES 89)
+
 # ------------- configure rapids-cmake --------------#
 
 include(../cmake/thirdparty/fetch_rapids.cmake)