注意

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

矩阵排序#

Argmax#

#include <raft/matrix/argmax.cuh>

namespace raft::matrix

template<typename math_t, typename idx_t, typename matrix_idx_t>
void argmax(raft::resources const &handle, raft::device_matrix_view<const math_t, matrix_idx_t, row_major> in, raft::device_vector_view<idx_t, matrix_idx_t> out)#

Argmax: 找到每行的最大值对应的列索引。

参数:

  • handle – [in] raft 句柄

  • in – [in] 大小为 (n_rows, n_cols) 的输入矩阵

  • out – [out] 大小为 n_rows 的输出向量

Argmin#

#include <raft/matrix/argmin.cuh>

namespace raft::matrix

template<typename math_t, typename idx_t, typename matrix_idx_t>
void argmin(raft::resources const &handle, raft::device_matrix_view<const math_t, matrix_idx_t, row_major> in, raft::device_vector_view<idx_t, matrix_idx_t> out)#

Argmin: 找到每行的最小值对应的列索引。

参数:

  • handle – [in] raft 句柄

  • in – [in] 大小为 (n_rows, n_cols) 的输入矩阵

  • out – [out] 大小为 n_rows 的输出向量

Select-K#

#include <raft/matrix/select_k.cuh>

namespace raft::matrix

enum class SelectAlgo : uint8_t#

用于选择 k 个最大邻居的算法。

关于 RAFT 中 select-k 算法工作原理的详细信息可以在论文“GPU 上的并行 Top-K 算法：一项全面研究和新方法”https://doi.org/10.1145/3581784.3607062 中找到。下面的 kRadix* 变体对应于论文中描述的“Air Top-k”算法，而 kWarp* 变体对应于“GridSelect”算法。

值

enumerator kAuto#

根据输入维度和 k 值自动选择 select-k 算法

enumerator kRadix8bits#

每次通过使用 8 位的基数选择

enumerator kRadix11bits#

每次通过使用 11 位的基数选择，融合最后一个过滤步骤

enumerator kRadix11bitsExtraPass#

每次通过使用 11 位的基数选择，不融合最后一个过滤步骤

enumerator kWarpAuto#

根据输入大小在 kWarpImmediate 和 kWarpFiltered 算法之间自动切换

enumerator kWarpImmediate#

此版本的 warp_sort 将每个输入元素添加到中间排序缓冲区中，因此每 Capacity 个输入元素执行一次排序步骤。

对于非常小的 len 值，首选此实现。

enumerator kWarpFiltered#

此版本的 warp_sort 在将每个输入元素添加到中间排序缓冲区之前，将其与当前估计的第 k 个值进行比较。这使得该算法在处理长输入序列时减少了排序步骤，但代价是每一步都需要进行额外的检查。

对于较大的 len 值，首选此实现。

enumerator kWarpDistributed#

此版本的 warp_sort 在将每个输入元素添加到中间排序缓冲区之前，将其与当前估计的第 k 个值进行比较。与 warp_sort_filtered 不同，它为 warp 中的所有线程（独立于子 warp 大小）保留一个分布式缓冲区，这使得刷新频率降低。

enumerator kWarpDistributedShm#

与 warp_sort_distributed 相同，但将临时值和索引缓冲区保存在给定的外部指针中（通常应传入一个共享内存指针）。

template<typename T, typename IdxT>
void select_k(raft::resources const &handle, raft::device_matrix_view<const T, int64_t, row_major> in_val, std::optional<raft::device_matrix_view<const IdxT, int64_t, row_major>> in_idx, raft::device_matrix_view<T, int64_t, row_major> out_val, raft::device_matrix_view<IdxT, int64_t, row_major> out_idx, bool select_min, bool sorted = false, SelectAlgo algo = SelectAlgo::kAuto)#

从输入数据的每行中选择 k 个最小或最大的键/值。

如果您将输入数据 in_val 视为一个具有 len 列和 batch_size 行的行主序矩阵，则此函数将选择每行中 k 个最小/最大值，并填充大小为 (batch_size, k) 的行主序矩阵 out_val。如果一行中的总值数小于 K，则 out_val 相应行中的额外位置将保留原始值。这同样适用于 out_idx。

使用示例

using namespace raft;
// get a 2D row-major array of values to search through
auto in_values = {... input device_matrix_view<const float, int64_t, row_major> ...}
// prepare output arrays
auto out_extents = make_extents<int64_t>(in_values.extent(0), k);
auto out_values  = make_device_mdarray<float>(handle, out_extents);
auto out_indices = make_device_mdarray<int64_t>(handle, out_extents);
// search `k` smallest values in each row
matrix::select_k<float, int64_t>(
  handle, in_values, std::nullopt, out_values.view(), out_indices.view(), true);

模板参数:

  • T – 键的类型（用于比较的值）。

  • IdxT – 索引类型（与键一起选择的值）。

参数:

  • handle – [in] 管理可重用资源的容器。

  • in_val – [in] 输入值 [batch_size, len]；这些值用于比较和选择。

  • in_idx – [in] 可选的输入载荷 [batch_size, len]；通常是对应 in_val 的索引。如果 in_idx 是 std::nullopt，则表示使用连续数组 0...len-1。

  • out_val – [out] 输出值 [batch_size, k]；in_val 每行的 k 个最小/最大值。

  • out_idx – [out] 输出载荷（例如索引）[batch_size, k]；与 out_val 一起选择的载荷。

  • select_min – [in] 标志，指示是否选择 k 个最小（true）或最大（false）的键。

  • sorted – [in] 是否确保选定的对按值排序

  • algo – [in] 要使用的选择算法

inline std::ostream &operator<<(std::ostream &os, const SelectAlgo &algo)#
template<typename T, typename IdxT>
void select_k(raft::resources const &handle, raft::device_csr_matrix_view<const T, IdxT, IdxT, IdxT> in_val, std::optional<raft::device_vector_view<const IdxT, IdxT>> in_idx, raft::device_matrix_view<T, IdxT, raft::row_major> out_val, raft::device_matrix_view<IdxT, IdxT, raft::row_major> out_idx, bool select_min, bool sorted = false, SelectAlgo algo = SelectAlgo::kAuto)#

从输入矩阵的每行中选择 k 个最小或最大的键/值。

此函数对一个行主序矩阵 in_val 进行操作，该矩阵的维度为 batch_size x len，从每行中选择 k 个最小或最大元素。然后将选定的元素存储在维度为 batch_size x k 的行主序输出矩阵 out_val 中。如果一行中的总值数小于 K，则 out_val 相应行中的额外位置将保留原始值。这同样适用于 out_idx。

模板参数:

  • T – 正在比较的元素类型（键）。

  • IdxT – 与键关联的索引类型。

参数:

  • handle – [in] 管理可重用资源的容器。

  • in_val – [in] CSR 格式的输入矩阵，逻辑稠密形状为 [batch_size, len]，包含要比较和选择的元素。

  • in_idx – [in] 与 in_val.values 关联的可选输入索引 [in_val.nnz]。如果 in_idx 是 std::nullopt，则默认为从 0 到 len-1 的连续数组。

  • out_val – [out] 输出矩阵 [in_val.get_n_row(), k]，存储从 in_val 每行中选择的 k 个最小/最大元素。

  • out_idx – [out] 输出索引 [in_val.get_n_row(), k]，对应于 out_val 中选定的元素。

  • select_min – [in] 标志，指示是否选择 k 个最小（true）或最大（false）元素。

  • sorted – [in] 是否确保选定的对按值排序

  • algo – [in] 要使用的选择算法

按列排序#

#include <raft/matrix/col_wise_sort.cuh>

namespace raft::matrix

template<typename in_t, typename out_t, typename matrix_idx_t, typename sorted_keys_t>
void sort_cols_per_row(raft::resources const &handle, raft::device_matrix_view<const in_t, matrix_idx_t, raft::row_major> in, raft::device_matrix_view<out_t, matrix_idx_t, raft::row_major> out, sorted_keys_t &&sorted_keys_opt)#

对行主序输入矩阵的每行中的列进行排序，并返回排序后的索引，这些索引被建模为键值对排序，其中键是输入矩阵，值是值的索引

模板参数:

  • in_t – 输入矩阵的元素类型

  • out_t – 输出矩阵的元素类型

  • matrix_idx_t – 用于矩阵索引的整数类型

  • sorted_keys_t – std::optional<raft::device_matrix_view<in_t, matrix_idx_t, raft::row_major>> sorted_keys_opt

参数:

  • handle – [in] raft 句柄

  • in – [in] 输入矩阵

  • out – [out] 输出值（索引）矩阵

  • sorted_keys_opt – [out] std::optional，排序键（输入）的输出矩阵

template<typename ...Args, typename = std::enable_if_t<sizeof...(Args) == 3>>
void sort_cols_per_row(Args... args)#

sort_keys_per_row 的重载，用于帮助编译器找到上述重载，以防用户为可选参数中的一个或两个传入 std::nullopt。

有关 sort_keys_per_row 的文档，请参见上文。