互操作性#
DLPack (C)#
近似最近邻(ANN)索引提供了一个通过C API构建和搜索索引的接口。[DLPack v0.8](dmlc/dlpack),一个张量接口框架,被用作与我们的C API交互的标准。
使用DLPack表示张量很简单,因为它是一个POD结构体,在运行时存储有关张量的信息。目前,来自DLPack v0.8的DLManagedTensor与我们的C API兼容,但我们很快将升级到来自DLPack v1.0的DLManagedTensorVersioned,因为它将帮助我们维护ABI和API的兼容性。
以下是如何使用DLManagedTensor表示设备内存的示例:
#include <dlpack/dlpack.h>
// Create data representation in host memory
float dataset[2][1] = {{0.2, 0.1}};
// copy data to device memory
float *dataset_dev;
cuvsRMMAlloc(&dataset_dev, sizeof(float) * 2 * 1);
cudaMemcpy(dataset_dev, dataset, sizeof(float) * 2 * 1, cudaMemcpyDefault);
// Use DLPack for representing the data as a tensor
DLManagedTensor dataset_tensor;
dataset_tensor.dl_tensor.data = dataset;
dataset_tensor.dl_tensor.device.device_type = kDLCUDA;
dataset_tensor.dl_tensor.ndim = 2;
dataset_tensor.dl_tensor.dtype.code = kDLFloat;
dataset_tensor.dl_tensor.dtype.bits = 32;
dataset_tensor.dl_tensor.dtype.lanes = 1;
int64_t dataset_shape[2] = {2, 1};
dataset_tensor.dl_tensor.shape = dataset_shape;
dataset_tensor.dl_tensor.strides = nullptr;
// free memory after use
cuvsRMMFree(dataset_dev);
请参考 cuVS C API 文档 以了解更多信息。
多维跨度 (C++)#
cuVS 构建在 RAFT 库中的 GPU 加速机器学习和数据挖掘原语之上。cuVS 中的大多数 C++ API 接受 mdspan 多维数组视图,用于表示高维数据,类似于 Numpy Python 库中的 ndarray。RAFT 还包含相应的拥有 mdarray 结构,它简化了主机和设备(GPU)内存中多维数据的分配和管理。
mdarray 是一个拥有所有权的对象,它在 RMM 上形成了一个便利层,并且可以在 RAFT 中使用多种不同的辅助函数来构造:
#include <raft/core/device_mdarray.hpp>
int n_rows = 10;
int n_cols = 10;
auto scalar = raft::make_device_scalar<float>(handle, 1.0);
auto vector = raft::make_device_vector<float>(handle, n_cols);
auto matrix = raft::make_device_matrix<float>(handle, n_rows, n_cols);
mdspan 是一个轻量级的非拥有视图,可以包装任何指针,维护用于访问元素的形状、布局和索引信息。
我们可以直接从上面的mdarray实例构造mdspan实例:
// Scalar mdspan on device
auto scalar_view = scalar.view();
// Vector mdspan on device
auto vector_view = vector.view();
// Matrix mdspan on device
auto matrix_view = matrix.view();
由于mdspan只是一个轻量级的包装器,我们也可以从上述mdarray实例中的底层数据句柄构造它。我们使用extent来获取有关mdarray或mdspan形状的信息。
#include <raft/core/device_mdspan.hpp>
auto scalar_view = raft::make_device_scalar_view(scalar.data_handle());
auto vector_view = raft::make_device_vector_view(vector.data_handle(), vector.extent(0));
auto matrix_view = raft::make_device_matrix_view(matrix.data_handle(), matrix.extent(0), matrix.extent(1));
当然,RAFT的mdspan/mdarray API不仅限于device。你也可以创建host变体:
#include <raft/core/host_mdarray.hpp>
#include <raft/core/host_mdspan.hpp>
int n_rows = 10;
int n_cols = 10;
auto scalar = raft::make_host_scalar<float>(handle, 1.0);
auto vector = raft::make_host_vector<float>(handle, n_cols);
auto matrix = raft::make_host_matrix<float>(handle, n_rows, n_cols);
auto scalar_view = raft::make_host_scalar_view(scalar.data_handle());
auto vector_view = raft::make_host_vector_view(vector.data_handle(), vector.extent(0));
auto matrix_view = raft::make_host_matrix_view(matrix.data_handle(), matrix.extent(0), matrix.extent(1));
请参考 RAFT 的 mdspan 文档 以了解更多信息。