Attention

The vector search and clustering algorithms in RAFT are being migrated to a new library dedicated to vector search called cuVS. We will continue to support the vector search algorithms in RAFT during this move, but will no longer update them after the RAPIDS 24.06 (June) release. We plan to complete the migration by RAPIDS 24.08 (August) release.

Utilities

RAFT contains numerous utility functions and primitives that are easily usable. This page provides C++ API references for the publicly-exposed utility functions.

Memory Pool

#include <raft/utils/memory_pool.cuh>

namespace raft

std::unique_ptr<rmm::mr::device_memory_resource> get_pool_memory_resource(rmm::mr::device_memory_resource *&mr, size_t initial_size)

Get a pointer to a pooled memory resource within the scope of the lifetime of the returned unique pointer.

This function is useful in the code where multiple repeated allocations/deallocations are expected. Use case example:

void my_func(..., size_t n, rmm::mr::device_memory_resource* mr = nullptr) {
  auto pool_guard = raft::get_pool_memory_resource(mr, 2 * n * sizeof(float));
  if (pool_guard){
    RAFT_LOG_INFO("Created a pool");
  } else {
    RAFT_LOG_INFO("Using the current default or explicitly passed device memory resource");
  }
  rmm::device_uvector<float> x(n, stream, mr);
  rmm::device_uvector<float> y(n, stream, mr);
  ...
}

Here, the new memory resource would be created within the function scope if the passed mr is null and the default resource is not a pool. After the call, mr contains a valid memory resource in any case.

Parameters:

• mr – [inout] if not null do nothing; otherwise get the current device resource and wrap it into a pool_memory_resource if necessary and return the pointer to the result.

• initial_size – if a new memory pool is created, this would be its initial size (rounded up to 256 bytes).

Returns:

if a new memory pool is created, it returns a unique_ptr to it; this managed pointer controls the lifetime of the created memory resource.