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.