注意

RAFT 中的向量搜索和聚类算法正在迁移到专门用于向量搜索的新库 cuVS。在此迁移过程中，我们将继续支持 RAFT 中的向量搜索算法，但在 RAPIDS 24.06（6 月）版本之后将不再更新它们。我们计划在 RAPIDS 24.10（10 月）版本之前完成迁移，并在 24.12（12 月）版本中将它们从 RAFT 中完全移除。

位集#

#include <raft/core/bitset.cuh>

namespace raft::core

template<typename bitset_t = uint32_t, typename index_t = uint32_t>
struct bitset_view#
#include <bitset.hpp>

RAFT 位集的视图。

这个轻量级结构存储了指向设备内存中位集及其长度的指针。它提供了一个 test() 设备函数来检查位集中给定索引是否已设置。

模板参数：

  • bitset_t – 位集数组的基础类型。默认为 uint32_t。

  • index_t – 使用的索引类型。默认为 uint32_t。

公有函数

inline _RAFT_HOST_DEVICE bitset_view(bitset_t *bitset_ptr, index_t bitset_len, index_t original_nbits = 0)#

从指向位集的设备指针创建位集视图。

参数：

  • bitset_ptr – 指向位集的设备指针

  • bitset_len – 位集中的位数

  • original_nbits – 位集创建时使用的原始位数，用于处理数据类型潜在的不匹配。这对于在使用 ANN 索引时非常有用，当位集最初创建时使用了与 cuVS ANN 索引当前支持的数据类型不同的类型。

inline _RAFT_HOST_DEVICE bitset_view(raft::device_vector_view<bitset_t, index_t> bitset_span, index_t bitset_len, index_t original_nbits = 0)#

从位集的设备向量视图创建位集视图。

参数：

  • bitset_span – 位集的设备向量视图

  • bitset_len – 位集中的位数

  • original_nbits – 位集创建时使用的原始位数，用于处理数据类型潜在的不匹配。这对于在使用 ANN 索引时非常有用，当位集最初创建时使用了与 cuVS ANN 索引当前支持的数据类型不同的类型。

inline _RAFT_HOST_DEVICE auto test(const index_t sample_index) const -> bool#

设备函数，用于测试位集中给定索引是否已设置。

参数：

sample_index – 要测试的单个索引

返回：

bool 如果索引在位集中未被取消设置，则为 True

inline _RAFT_HOST_DEVICE auto operator[](const index_t sample_index) const -> bool#

设备函数，用于测试位集中给定索引是否已设置。

参数：

sample_index – 要测试的单个索引

返回：

bool 如果索引在位集中未被取消设置，则为 True

inline _RAFT_HOST_DEVICE void set (const index_t sample_index, bool set_value) const

设备函数，用于将位集中给定索引设置为 set_value。

参数：

  • sample_index – 要设置的索引

  • set_value – 要将位设置为什么值（true 或 false）

inline _RAFT_HOST_DEVICE auto data() -> bitset_t*#

获取指向位集的设备指针。

inline _RAFT_HOST_DEVICE auto size() const -> index_t#

获取位集表示的位数。

inline _RAFT_HOST_DEVICE auto n_elements() const -> index_t#

获取位集表示使用的元素数量。

void count(const raft::resources &res, raft::device_scalar_view<index_t> count_gpu_scalar) const#

返回 count_gpu_scalar 中设置为 true 的位数。

参数：

  • res – [in] RAFT 资源

  • count_gpu_scalar – [out] 用于存储计数的设备标量

inline index_t count(const raft::resources &res) const#

返回设置为 true 的位数。

参数：

res – RAFT 资源

返回：

index_t 设置为 true 的位数

void repeat(const raft::resources &res, index_t times, bitset_t *output_device_ptr) const#

重复位集数据并将其复制到输出设备指针。

此函数获取存储在设备内存中的原始位集数据，并将其重复指定次数到设备内存中的新位置。逐位复制，以确保即使位数 (bitset_len_) 不是位集元素大小（例如 uint32_t 为 32 位）的倍数，这些位也能紧密打包，行之间没有间隙。

调用者必须确保输出设备指针已分配足够的内存，以容纳 times * bitset_len 位，其中 bitset_len 是原始位集中的位数。此函数使用 Thrust 并行算法在 GPU 上高效执行操作。

参数：

  • res – 用于管理 CUDA 流和执行策略的 RAFT 资源。

  • times – 位集数据在输出中应重复的次数。

  • output_device_ptr – 存储重复位集数据的设备指针。

double sparsity(const raft::resources &res) const#

计算位集的稀疏度（0 的分数）。

此函数计算位集的稀疏度，稀疏度定义为未设置位（0）与位集中总位数之比。如果总位数为零，则函数返回 1.0，表示该集合完全稀疏。

此 API 将在 res 的流上同步。

参数：

res – 用于管理 CUDA 流和执行策略的 RAFT 资源。

返回：

double 位集的稀疏度，即未设置位的比例。

inline index_t get_original_nbits() const#

获取位集的原始位数。

template<typename csr_matrix_t>
void to_csr(const raft::resources &res, csr_matrix_t &csr) const#

转换为压缩稀疏行 (CSR) 格式矩阵。

此方法将位集视图转换为 CSR 矩阵表示，其中位集中的每个“1”位对应于 CSR 矩阵中的一个非零项。位集格式仅支持单行矩阵，因此如果 CSR 矩阵需要多行，则位集数据在输出中为每一行重复。

示例用法

#include <raft/core/resource/cuda_stream.hpp>
#include <raft/sparse/convert/csr.cuh>
#include <rmm/device_uvector.hpp>

using bitset_t = uint32_t;
using index_t  = int;
using value_t  = float;

raft::resources handle;
auto stream    = resource::get_cuda_stream(handle);
index_t n_rows = 3;
index_t n_cols = 30;

// Compute bitset size and initialize device memory
index_t bitset_size = (n_cols + sizeof(bitset_t) * 8 - 1) / (sizeof(bitset_t) * 8);
rmm::device_uvector<bitset_t> bitset_d(bitset_size, stream);
std::vector<bitset_t> bitset_h = {
  bitset_t(0b11001010),
};  // Example bitset, with 4 non-zero entries.

raft::copy(bitset_d.data(), bitset_h.data(), bitset_h.size(), stream);

// Create bitset view and CSR matrix
auto bitset_view = raft::core::bitset_view<bitset_t, index_t>(bitset_d.data(), n_cols);
auto csr = raft::make_device_csr_matrix<value_t, index_t>(handle, n_rows, n_cols, 4 * n_rows);

// Convert bitset to CSR
bitset_view.to_csr(handle, csr);
resource::sync_stream(handle);

// Results:
// csr.indptr  = [0, 4, 8, 12];
// csr.indices = [1, 3, 6, 7,
//                1, 3, 6, 7,
//                1, 3, 6, 7];
// csr.values  = [1, 1, 1, 1,
//                1, 1, 1, 1,
//                1, 1, 1, 1];

调用者必须确保：csr 矩阵已预分配维度和非零计数与预期输出匹配，即 nnz_for_csr = nnz_for_bitset * n_rows。

模板参数：

csr_matrix_t – 指定 CSR 矩阵类型，约束为 raft::device_csr_matrix。

参数：

  • res – [in] 用于管理 CUDA 流和执行策略的 RAFT 资源。

  • csr – [out] 存储结果 CSR 矩阵的输出参数。位集中的每个“1”位对应于 CSR 矩阵中的一个非零元素。

公有静态函数

static inline size_t eval_n_elements(size_t bitset_len)#

计算存储位集所需的 bitset_t 元素数量。

此函数计算存储位集所需的 bitset_t 元素数量，确保所有位都被考虑在内。如果位集长度不是 bitset_t 大小（以位为单位）的倍数，则计算会向上取整以将剩余位包含在一个额外的 bitset_t 元素中。

参数：

bitset_len – 位集的总长度（以位为单位）。

返回：

size_t 存储位集所需的 bitset_t 元素数量。

template<typename bitset_t = uint32_t, typename index_t = uint32_t>
struct bitset#
#include <bitset.hpp>

RAFT 位集。

此结构封装了设备内存中的位集。它提供了一个 view() 方法来获取位集的设备可用轻量级视图。每个索引由位集中的一个单独位表示。使用的总字节数为 ceil(bitset_len / 8)。

模板参数：

  • bitset_t – 位集数组的基础类型。默认为 uint32_t。

  • index_t – 使用的索引类型。默认为 uint32_t。

公有函数

bitset(const raft::resources &res, raft::device_vector_view<const index_t, index_t> mask_index, index_t bitset_len, bool default_value = true)#

构造一个新的位集对象，包含要取消设置的索引列表。

参数：

  • res – RAFT 资源

  • mask_index – 要在位集中取消设置的索引列表

  • bitset_len – 位集的长度

  • default_value – 要将位设置为什么值的默认值。默认为 true。

bitset(const raft::resources &res, index_t bitset_len, bool default_value = true)#

构造一个新的位集对象。

参数：

  • res – RAFT 资源

  • bitset_len – 位集的长度

  • default_value – 要将位设置为什么值的默认值。默认为 true。

inline raft::core::bitset_view<bitset_t, index_t> view()#

创建一个位集的设备可用视图。

返回：

bitset_view<bitset_t, index_t>

inline bitset_t *data()#

获取指向位集的设备指针。

inline index_t size() const#

获取位集表示的位数。

inline index_t n_elements() const#

获取位集表示使用的元素数量。

inline raft::device_vector_view<bitset_t, index_t> to_mdspan()#

获取当前位集的 mdspan 视图。

void resize(const raft::resources &res, index_t new_bitset_len, bool default_value = true)#

调整位集大小。如果请求的大小更大，将分配新内存并设置为默认值。

参数：

  • res – RAFT 资源

  • new_bitset_len – 位集的新大小

  • default_value – 用于初始化新位的默认值

template<typename output_t = bool>
void test(const raft::resources &res, raft::device_vector_view<const index_t, index_t> queries, raft::device_vector_view<output_t, index_t> output) const#

测试位集中的索引列表。

模板参数：

output_t – 测试的输出类型。默认为 bool。

参数：

  • res – RAFT 资源

  • queries – 要测试的索引列表

  • output – 输出列表

void set(const raft::resources &res, raft::device_vector_view<const index_t, index_t> mask_index, bool set_value = false)#

将位集中的索引列表设置为 set_value。

参数：

  • res – RAFT 资源

  • mask_index – 要从位集中移除的索引

  • set_value – 要将位设置为什么值（true 或 false）

void flip(const raft::resources &res)#

翻转位集中的所有位。

参数：

res – RAFT 资源

void reset(const raft::resources &res, bool default_value = true)#

重置位集中的位。

参数：

  • res – RAFT 资源

  • default_value – 要将位设置为什么值（true 或 false）

void count(const raft::resources &res, raft::device_scalar_view<index_t> count_gpu_scalar)#

返回 count_gpu_scalar 中设置为 true 的位数。

参数：

  • res – [in] RAFT 资源

  • count_gpu_scalar – [out] 用于存储计数的设备标量

inline index_t count(const raft::resources &res)#

返回设置为 true 的位数。

参数：

res – RAFT 资源

返回：

index_t 设置为 true 的位数

inline bool any(const raft::resources &res)#

检查位集中是否有任何位设置为 true。

参数：

res – RAFT 资源

inline bool all(const raft::resources &res)#

检查位集中是否所有位都设置为 true。

参数：

res – RAFT 资源

inline bool none(const raft::resources &res)#

检查位集中是否没有位设置为 true。

参数：

res – RAFT 资源