cuVS API 基础#

内存管理#

集中式内存管理允许灵活配置分配策略,例如跨库边界共享相同的CUDA内存池。cuVS使用[RMM](rapidsai/rmm)库,这减轻了在GPU加速库中全局配置不同分配策略的负担。

RMM 目前有适用于 C++ 和 Python 的 API。

C++#

以下是一个在C++中配置RMM以使用池分配器的示例(源自RMM示例这里):

rmm::mr::cuda_memory_resource cuda_mr;
// Construct a resource that uses a coalescing best-fit pool allocator
// With the pool initially half of available device memory
auto initial_size = rmm::percent_of_free_device_memory(50);
rmm::mr::pool_memory_resource<rmm::mr::cuda_memory_resource> pool_mr{&cuda_mr, initial_size};
rmm::mr::set_current_device_resource(&pool_mr); // Updates the current device resource pointer to `pool_mr`
rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource(); // Points to `pool_mr`

Python#

以及Python中的相应代码(源自RMM示例这里):

import rmm
pool = rmm.mr.PoolMemoryResource(
  rmm.mr.CudaMemoryResource(),
  initial_pool_size=2**30,
  maximum_pool_size=2**32)
rmm.mr.set_current_device_resource(pool)

资源管理#

cuVS 使用来自 RAFT 库的 API,该库包含机器学习和数据挖掘的基本操作,用于集中和重用昂贵的资源,例如内存管理。以下代码示例展示了如何创建这些资源以在本指南中使用。

请参阅 RAFT 的 资源 API 文档 以获取更多信息。

C#

#include <cuda_runtime.h>
#include <cuvs/core/c_api.h>

cuvsResources_t res;
cuvsResourcesCreate(&res);

// ... do some processing ...

cuvsResourcesDestroy(res);

C++#

#include <raft/core/device_resources.hpp>

raft::device_resources res;

Python#

import pylibraft

res = pylibraft.common.DeviceResources()

Rust#

let res = cuvs::Resources::new()?;