快速入门

本指南旨在为与 RAFT 的 C++ 和 Python API 交互提供快速入门教程。

RAPIDS Memory Manager (RMM)

RAFT 严重依赖于 RMM 库，该库简化了在使用它的库中全局配置不同分配策略的负担。

多维 Span 和 Array

RAFT 中的大多数 API 都接受 mdspan 多维数组视图，用于表示高维数据，类似于 Numpy Python 库中的 ndarray。RAFT 还包含相应的拥有型 mdarray 结构，它简化了 host 和 device (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));

以及 managed 变体。

#include <raft/core/device_mdspan.hpp>

int n_rows = 10;
int n_cols = 10;

auto matrix = raft::make_managed_mdspan(managed_ptr, raft::make_matrix_extents(n_rows, n_cols));

您还可以创建带步长 (strided) 的 mdspan。

#include <raft/core/device_mdspan.hpp>

int n_elements = 10;
int stride = 10;

auto vector = raft::make_device_vector_view(vector_ptr, raft::make_vector_strided_layout(n_elements, stride));

C++ 示例

RAFT 中的大多数 primitives 都接受 raft::handle_t 对象来管理创建成本高昂的资源，例如 CUDA streams、stream pools，以及指向其他 CUDA 库（如 cublas 和 cusolver）的句柄。

下面的示例演示了如何创建 RAFT 句柄，并将其与 device_matrix 和 device_vector 一起使用，以分配内存、生成随机簇，并使用 NVIDIA cuVS 计算成对欧几里得距离。

#include <raft/core/handle.hpp>
#include <raft/core/device_mdarray.hpp>
#include <raft/random/make_blobs.cuh>
#include <cuvs/distance/distance.hpp>

raft::handle_t handle;

int n_samples = 5000;
int n_features = 50;

auto input = raft::make_device_matrix<float>(handle, n_samples, n_features);
auto labels = raft::make_device_vector<int>(handle, n_samples);
auto output = raft::make_device_matrix<float>(handle, n_samples, n_samples);

raft::random::make_blobs(handle, input.view(), labels.view());

auto metric = cuvs::distance::DistanceType::L2SqrtExpanded;
cuvs::distance::pairwise_distance(handle, input.view(), input.view(), output.view(), metric);

Python 示例

pylibraft 包包含用于 RAFT 算法和 primitives 的 Python API。pylibraft 重量轻，依赖少，可以接受任何支持 __cuda_array_interface__ 的对象，例如 CuPy 的 ndarray，因此能很好地集成到其他库中。该包中公开的 RAFT 算法数量在每个版本中都在持续增长。

下面的示例演示了使用 NVIDIA cuVS 库计算 CuPy 数组之间的成对欧几里得距离。请注意，CuPy 不是 pylibraft 的必需依赖项。

import cupy as cp

from cuvs.distance import pairwise_distance

n_samples = 5000
n_features = 50

in1 = cp.random.random_sample((n_samples, n_features), dtype=cp.float32)
in2 = cp.random.random_sample((n_samples, n_features), dtype=cp.float32)

output = pairwise_distance(in1, in2, metric="euclidean")

上述示例中的 output 数组类型为 raft.common.device_ndarray，它支持 cuda_array_interface，使其与同样支持它的其他库（如 CuPy、Numba 和 PyTorch）具有互操作性。CuPy 支持 DLPack，这也使得从 raft.common.device_ndarray 到 JAX 和 Tensorflow 的零拷贝转换成为可能。

下面是将输出 pylibraft.common.device_ndarray 转换为 CuPy 数组的示例

cupy_array = cp.asarray(output)

以及转换为 PyTorch tensor

import torch

torch_tensor = torch.as_tensor(output, device='cuda')

当相应的库已安装并在您的环境中可用时，所有 RAFT 计算 API 也可以通过设置全局配置选项来自动完成此转换

import pylibraft.config
pylibraft.config.set_output_as("cupy")  # All compute APIs will return cupy arrays
pylibraft.config.set_output_as("torch") # All compute APIs will return torch tensors

您还可以指定一个接受 pylibraft.common.device_ndarray 并执行自定义转换的 callable。以下示例将所有输出转换为 numpy 数组

pylibraft.config.set_output_as(lambda device_ndarray: return device_ndarray.copy_to_host())

pylibraft 还支持写入预分配的输出数组，因此任何支持 __cuda_array_interface__ 的数组都可以就地写入。

import cupy as cp

from cuvs.distance import pairwise_distance

n_samples = 5000
n_features = 50

in1 = cp.random.random_sample((n_samples, n_features), dtype=cp.float32)
in2 = cp.random.random_sample((n_samples, n_features), dtype=cp.float32)
output = cp.empty((n_samples, n_samples), dtype=cp.float32)

pairwise_distance(in1, in2, out=output, metric="euclidean")