注意

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

映射与归约#

合并归约#

#include <raft/linalg/coalesced_reduction.cuh>

namespace raft::linalg

template<typename InValueType, typename LayoutPolicy, typename OutValueType, typename IdxType, typename MainLambda = raft::identity_op, typename ReduceLambda = raft::add_op, typename FinalLambda = raft::identity_op>
void coalesced_reduction(raft::resources const &handle, raft::device_matrix_view<const InValueType, IdxType, LayoutPolicy> data, raft::device_vector_view<OutValueType, IdxType> dots, OutValueType init, bool inplace = false, MainLambda main_op = raft::identity_op(), ReduceLambda reduce_op = raft::add_op(), FinalLambda final_op = raft::identity_op())#

沿着主维度计算输入矩阵的归约。当期望的归约方向与内存布局维度一致时，应使用此 API。例如，行主序矩阵将沿着列进行归约，而列主序矩阵将沿着行进行归约。

模板参数:

  • InValueType – 底层 raft::matrix_view 的输入数据类型

  • LayoutPolicy – 输入/输出的布局（行主序或列主序）

  • OutValueType – 底层 raft::matrix_view 和归约的输出数据类型

  • IndexType – 用于寻址的整数类型

  • MainLambda – 累加时应用的一元 lambda 函数（例如：L1 或 L2 范数）。它必须是一个支持以下输入和输出的“可调用对象”。

  • ReduceLambda – 用于归约的二元 lambda 函数（例如：L2 范数的加法（+））。它必须是一个支持以下输入和输出的“可调用对象”。

  • FinalLambda – 在 STG 前应用的最终 lambda 函数（例如：L2 范数的开方）。它必须是一个支持以下输入和输出的“可调用对象”。

参数:

  • handle – raft::resources

  • data – [in] 类型为 raft::device_matrix_view 的输入

  • dots – [out] 类型为 raft::device_matrix_view 的输出

  • init – [in] 用于归约的初始值

  • inplace – [in] 归约结果是在原位添加还是覆盖旧值？

  • main_op – [in] 在归约前应用的融合逐元素运算

  • reduce_op – [in] 融合的二元归约运算

  • final_op – [in] 在存储结果前应用的融合逐元素运算

映射#

#include <raft/linalg/map.cuh>

namespace raft::linalg

template<typename OutType, typename Func, typename ...InTypes, typename = raft::enable_if_output_device_mdspan<OutType>, typename = raft::enable_if_input_device_mdspan<InTypes...>>
void map(const raft::resources &res, OutType out, Func f, InTypes... ins)#

在零个或多个相同大小的输入 mdspan 上映射一个函数。

应用于 k 个输入的算法可以用以下伪代码描述

for (auto i: [0 ... out.size()]) {
  out[i] = f(in_0[i], in_1[i], ..., in_k[i])
}

性能说明：如果可能，此函数使用向量化 CUDA 加载/存储指令加载参数数组并存储输出数组。向量化的大小取决于最大输入/输出元素类型的大小以及所有指针的对齐方式。

用法示例

#include <raft/core/device_mdarray.hpp>
#include <raft/core/resources.hpp>
#include <raft/core/operators.hpp>
#include <raft/linalg/map.cuh>

auto input = raft::make_device_vector<int>(res, n);
... fill input ..
auto squares = raft::make_device_vector<int>(res, n);
raft::linalg::map_offset(res, squares.view(), raft::sq_op{}, input.view());

模板参数:

  • OutType – 结果的数据类型 (device_mdspan)

  • Func – 执行实际操作的设备 lambda 函数

  • InTypes – 输入的数据类型 (device_mdspan)

参数:

  • res – [in] raft::resources

  • out – [out] 映射操作的输出 (device_mdspan)

  • f – [in] 设备 lambda 函数 (InTypes::value_type xs…) -> OutType::value_type

  • ins – [in] 输入（每个的大小与输出相同）(device_mdspan)

template<typename InType1, typename OutType, typename Func, typename = raft::enable_if_output_device_mdspan<OutType>, typename = raft::enable_if_input_device_mdspan<InType1>>
void map(const raft::resources &res, InType1 in1, OutType out, Func f)#

在一个 mdspan 上映射一个函数。

模板参数:

  • InType1 – 输入的数据类型 (device_mdspan)

  • OutType – 结果的数据类型 (device_mdspan)

  • Func – 执行实际操作的设备 lambda 函数

参数:

  • res – [in] raft::resources

  • in1 – [in] 输入（与输出大小相同）(device_mdspan)

  • out – [out] 映射操作的输出 (device_mdspan)

  • f – [in] 设备 lambda 函数 (InType1::value_type x) -> OutType::value_type

template<typename InType1, typename InType2, typename OutType, typename Func, typename = raft::enable_if_output_device_mdspan<OutType>, typename = raft::enable_if_input_device_mdspan<InType1, InType2>>
void map(const raft::resources &res, InType1 in1, InType2 in2, OutType out, Func f)#

在两个 mdspan 上映射一个函数。

模板参数:

  • InType1 – 输入的数据类型 (device_mdspan)

  • InType2 – 输入的数据类型 (device_mdspan)

  • OutType – 结果的数据类型 (device_mdspan)

  • Func – 执行实际操作的设备 lambda 函数

参数:

  • res – [in] raft::resources

  • in1 – [in] 输入（与输出大小相同）(device_mdspan)

  • in2 – [in] 输入（与输出大小相同）(device_mdspan)

  • out – [out] 映射操作的输出 (device_mdspan)

  • f – [in] 设备 lambda 函数 (InType1::value_type x1, InType2::value_type x2) -> OutType::value_type

template<typename InType1, typename InType2, typename InType3, typename OutType, typename Func, typename = raft::enable_if_output_device_mdspan<OutType>, typename = raft::enable_if_input_device_mdspan<InType1, InType2, InType3>>
void map(const raft::resources &res, InType1 in1, InType2 in2, InType3 in3, OutType out, Func f)#

在三个 mdspan 上映射一个函数。

模板参数:

  • InType1 – 输入 1 的数据类型 (device_mdspan)

  • InType2 – 输入 2 的数据类型 (device_mdspan)

  • InType3 – 输入 3 的数据类型 (device_mdspan)

  • OutType – 结果的数据类型 (device_mdspan)

  • Func – 执行实际操作的设备 lambda 函数

参数:

  • res – [in] raft::resources

  • in1 – [in] 输入 1（与输出大小相同）(device_mdspan)

  • in2 – [in] 输入 2（与输出大小相同）(device_mdspan)

  • in3 – [in] 输入 3（与输出大小相同）(device_mdspan)

  • out – [out] 映射操作的输出 (device_mdspan)

  • f – [in] 设备 lambda 函数 (InType1::value_type x1, InType2::value_type x2, InType3::value_type x3) -> OutType::value_type

template<typename OutType, typename Func, typename ...InTypes, typename = raft::enable_if_output_device_mdspan<OutType>, typename = raft::enable_if_input_device_mdspan<InTypes...>>
void map_offset(const raft::resources &res, OutType out, Func f, InTypes... ins)#

在从零开始的平面索引（元素偏移量）和零个或多个输入上映射一个函数。

应用于 k 个输入的算法可以用以下伪代码描述

for (auto i: [0 ... out.size()]) {
  out[i] = f(i, in_0[i], in_1[i], ..., in_k[i])
}

性能说明：如果可能，此函数使用向量化 CUDA 加载/存储指令加载参数数组并存储输出数组。向量化的大小取决于最大输入/输出元素类型的大小以及所有指针的对齐方式。

用法示例

#include <raft/core/device_mdarray.hpp>
#include <raft/core/resources.hpp>
#include <raft/core/operators.hpp>
#include <raft/linalg/map.cuh>

auto squares = raft::make_device_vector<int>(handle, n);
raft::linalg::map_offset(res, squares.view(), raft::sq_op{});

模板参数:

  • OutType – 结果的数据类型 (device_mdspan)

  • Func – 执行实际操作的设备 lambda 函数

  • InTypes – 输入的数据类型 (device_mdspan)

参数:

  • res – [in] raft::resources

  • out – [out] 映射操作的输出 (device_mdspan)

  • f – [in] 设备 lambda 函数 (auto offset, InTypes::value_type xs…) -> OutType::value_type

  • ins – [in] 输入（每个的大小与输出相同）(device_mdspan)

template<typename InType1, typename OutType, typename Func, typename = raft::enable_if_output_device_mdspan<OutType>, typename = raft::enable_if_input_device_mdspan<InType1>>
void map_offset(const raft::resources &res, InType1 in1, OutType out, Func f)#

在从零开始的平面索引（元素偏移量）和一个 mdspan 上映射一个函数。

模板参数:

  • InType1 – 输入的数据类型 (device_mdspan)

  • OutType – 结果的数据类型 (device_mdspan)

  • Func – 执行实际操作的设备 lambda 函数

参数:

  • res – [in] raft::resources

  • in1 – [in] 输入（与输出大小相同）(device_mdspan)

  • out – [out] 映射操作的输出 (device_mdspan)

  • f – [in] 设备 lambda 函数 (auto offset, InType1::value_type x) -> OutType::value_type

template<typename InType1, typename InType2, typename OutType, typename Func, typename = raft::enable_if_output_device_mdspan<OutType>, typename = raft::enable_if_input_device_mdspan<InType1, InType2>>
void map_offset(const raft::resources &res, InType1 in1, InType2 in2, OutType out, Func f)#

在从零开始的平面索引（元素偏移量）和两个 mdspan 上映射一个函数。

模板参数:

  • InType1 – 输入的数据类型 (device_mdspan)

  • InType2 – 输入的数据类型 (device_mdspan)

  • OutType – 结果的数据类型 (device_mdspan)

  • Func – 执行实际操作的设备 lambda 函数

参数:

  • res – [in] raft::resources

  • in1 – [in] 输入（与输出大小相同）(device_mdspan)

  • in2 – [in] 输入（与输出大小相同）(device_mdspan)

  • out – [out] 映射操作的输出 (device_mdspan)

  • f – [in] 设备 lambda 函数 (auto offset, InType1::value_type x1, InType2::value_type x2) -> OutType::value_type

template<typename InType1, typename InType2, typename InType3, typename OutType, typename Func, typename = raft::enable_if_output_device_mdspan<OutType>, typename = raft::enable_if_input_device_mdspan<InType1, InType2, InType3>>
void map_offset(const raft::resources &res, InType1 in1, InType2 in2, InType3 in3, OutType out, Func f)#

在从零开始的平面索引（元素偏移量）和三个 mdspan 上映射一个函数。

模板参数:

  • InType1 – 输入 1 的数据类型 (device_mdspan)

  • InType2 – 输入 2 的数据类型 (device_mdspan)

  • InType3 – 输入 3 的数据类型 (device_mdspan)

  • OutType – 结果的数据类型 (device_mdspan)

  • Func – 执行实际操作的设备 lambda 函数

参数:

  • res – [in] raft::resources

  • in1 – [in] 输入 1（与输出大小相同）(device_mdspan)

  • in2 – [in] 输入 2（与输出大小相同）(device_mdspan)

  • in3 – [in] 输入 3（与输出大小相同）(device_mdspan)

  • out – [out] 映射操作的输出 (device_mdspan)

  • f – [in] 设备 lambda 函数 (auto offset, InType1::value_type x1, InType2::value_type x2, InType3::value_type x3) -> OutType::value_type

映射归约#

#include <raft/linalg/map_reduce.cuh>

namespace raft::linalg

template<typename InValueType, typename MapOp, typename ReduceLambda, typename IndexType, typename OutValueType, typename ScalarIdxType, typename ...Args>
void map_reduce(raft::resources const &handle, raft::device_vector_view<const InValueType, IndexType> in, raft::device_scalar_view<OutValueType, ScalarIdxType> out, OutValueType neutral, MapOp map, ReduceLambda op, Args... args)#

映射后进行通用归约操作的 CUDA 版本。

模板参数:

  • InValueType – 输入的数据类型

  • MapOp – 执行实际映射操作的设备 lambda 函数

  • ReduceLambda – 执行实际归约的设备 lambda 函数

  • IndexType – 索引类型

  • OutValueType – 输出的数据类型

  • ScalarIdxType – 标量的索引类型

  • Args – 附加参数

参数:

  • handle – [in] raft::resources

  • in – [in] 类型为 raft::device_vector_view 的输入

  • neutral – [in] 归约运算的中性元素。例如：求和为 0，乘法为 1，最小值取 +Inf，最大值取 -Inf

  • out – [out] 输出的归约值，假定为 raft::device_scalar_view 类型

  • map – [in] 融合的设备 lambda 函数

  • op – [in] 融合的归约设备 lambda 函数

  • args – [in] 附加输入数组

均方误差#

#include <raft/linalg/mean_squared_error.cuh>

namespace raft::linalg

template<typename InValueType, typename IndexType, typename OutValueType>
void mean_squared_error(raft::resources const &handle, raft::device_vector_view<const InValueType, IndexType> A, raft::device_vector_view<const InValueType, IndexType> B, raft::device_scalar_view<OutValueType, IndexType> out, OutValueType weight)#

均方误差函数 mean((A-B)**2) 的 CUDA 版本

模板参数:

  • InValueType – 输入数据类型

  • IndexType – 输入/输出索引类型

  • OutValueType – 输出数据类型

  • TPB – 每块线程数

参数:

范数#

#include <raft/linalg/norm.cuh>

namespace raft::linalg

template<typename ElementType, typename OutputType, typename LayoutPolicy, typename IndexType, typename Lambda = raft::identity_op>
void norm(raft::resources const &handle, raft::device_matrix_view<const ElementType, IndexType, LayoutPolicy> in, raft::device_vector_view<OutputType, IndexType> out, NormType type, Apply apply, Lambda fin_op = raft::identity_op())#

计算输入矩阵的范数并执行 fin_op。

模板参数:

  • ElementType – 输入数据类型

  • OutType – 输出数据类型

  • LayoutPolicy – 输入的布局（raft::row_major 或 raft::col_major）

  • IdxType – 用于寻址的整数类型

  • Lambda – 设备最终 lambda 函数

参数:

  • handle – [in] raft::resources

  • in – [in] 输入 raft::device_matrix_view

  • out – [out] 输出 raft::device_vector_view

  • type – [in] 应用的范数类型

  • apply – [in] 范数是应用于行方向（raft::linalg::Apply::ALONG_ROWS）还是列方向（raft::linalg::Apply::ALONG_COLUMNS）

  • fin_op – [in] 最终的 lambda 操作

归一化#

#include <raft/linalg/normalize.cuh>

namespace raft::linalg

template<typename ElementType, typename IndexType, typename MainLambda, typename ReduceLambda, typename FinalLambda>
void row_normalize(raft::resources const &handle, raft::device_matrix_view<const ElementType, IndexType, row_major> in, raft::device_matrix_view<ElementType, IndexType, row_major> out, ElementType init, MainLambda main_op, ReduceLambda reduce_op, FinalLambda fin_op, ElementType eps = ElementType(1e-8))#

根据由 main_op, reduce_op 和 fin_op 定义的范数来划分行。

模板参数:

  • ElementType – 输入/输出数据类型

  • IndexType – 用于寻址的整数类型

  • MainLambda – main_op 的类型

  • ReduceLambda – reduce_op 的类型

  • FinalLambda – fin_op 的类型

参数:

  • handle – [in] raft::resources

  • in – [in] 输入 raft::device_matrix_view

  • out – [out] 输出 raft::device_matrix_view

  • init – [in] 初始化值，即归约操作的单位元

  • main_op – [in] 在归约元素之前应用于元素的运算（例如 L2 范数的平方）

  • reduce_op – [in] 归约一对元素的运算（例如 L2 范数的求和）

  • fin_op – [in] 应用于归约结果一次以完成范数计算的运算（例如 L2 范数的平方根）

  • eps – [in] 如果范数低于 eps，则该行被视为零，不进行除法运算

template<typename ElementType, typename IndexType>
void row_normalize(raft::resources const &handle, raft::device_matrix_view<const ElementType, IndexType, row_major> in, raft::device_matrix_view<ElementType, IndexType, row_major> out, NormType norm_type, ElementType eps = ElementType(1e-8))#

根据范数来划分行。

模板参数:

  • ElementType – 输入/输出数据类型

  • IndexType – 用于寻址的整数类型

参数:

  • handle – [in] raft::resources

  • in – [in] 输入 raft::device_matrix_view

  • out – [out] 输出 raft::device_matrix_view

  • norm_type – [in] 要应用的范数类型

  • eps – [in] 如果范数低于 eps，则该行被视为零，不进行除法运算

归约#

#include <raft/linalg/reduce.cuh>

namespace raft::linalg

template<typename InElementType, typename LayoutPolicy, typename OutElementType = InElementType, typename IdxType = std::uint32_t, typename MainLambda = raft::identity_op, typename ReduceLambda = raft::add_op, typename FinalLambda = raft::identity_op>
void reduce(raft::resources const &handle, raft::device_matrix_view<const InElementType, IdxType, LayoutPolicy> data, raft::device_vector_view<OutElementType, IdxType> dots, OutElementType init, Apply apply, bool inplace = false, MainLambda main_op = raft::identity_op(), ReduceLambda reduce_op = raft::add_op(), FinalLambda final_op = raft::identity_op())#

沿指定维度计算输入矩阵的归约。该 API 计算底层存储为行主或列主矩阵的归约，同时允许选择归约维度。根据选择的归约维度，内存访问可能是合并的或带步长的。在加法归约的情况下，将执行补偿求和以减少数值误差。请注意，补偿不会等同于顺序补偿，以保持并行效率。

模板参数:

  • InElementType – 底层 raft::matrix_view 的输入数据类型

  • LayoutPolicy – 输入/输出的布局（行主序或列主序）

  • OutElementType – 底层 raft::matrix_view 和归约的输出数据类型

  • IndexType – 用于寻址的整数类型

  • MainLambda – 累加时应用的一元 lambda 函数（例如：L1 或 L2 范数）。它必须是一个支持以下输入和输出的“可调用对象”。

  • ReduceLambda – 用于归约的二元 lambda 函数（例如：L2 范数的加法（+））。它必须是一个支持以下输入和输出的“可调用对象”。

  • FinalLambda – 在 STG 前应用的最终 lambda 函数（例如：L2 范数的开方）。它必须是一个支持以下输入和输出的“可调用对象”。

参数:

  • handle – [in] raft::resources

  • data – [in] 类型为 raft::device_matrix_view 的输入

  • dots – [out] 类型为 raft::device_matrix_view 的输出

  • init – [in] 用于归约的初始值

  • apply – [in] 是沿行还是沿列归约（使用 raft::linalg::Apply）

  • main_op – [in] 在归约前应用的融合逐元素运算

  • reduce_op – [in] 融合的二元归约运算

  • final_op – [in] 在存储结果前应用的融合逐元素运算

  • inplace – [in] 归约结果是在原位添加还是覆盖旧值？

按键归约列#

#include <raft/linalg/reduce_cols_by_key.cuh>

namespace raft::linalg

template<typename ElementType, typename KeyType = ElementType, typename IndexType = std::uint32_t>
void reduce_cols_by_key(raft::resources const &handle, raft::device_matrix_view<const ElementType, IndexType, raft::row_major> data, raft::device_vector_view<const KeyType, IndexType, keys, raft::device_matrix_view<ElementType, IndexType, raft::row_major> out, IndexType nkeys = 0, bool reset_sums = true)#

计算给定键的矩阵列的求和归约。TODO: 支持通用归约 lambda rapidsai/raft#860。

模板参数:

  • ElementType – 输入数据类型（以及输出归约矩阵）

  • KeyType – 键的数据类型

  • IndexType – 索引算术类型

参数:

  • handle – [in] raft::resources

  • data – [in] 输入数据 (dim = nrows x ncols)。假定其为行主布局，类型为 raft::device_matrix_view

  • keys – [in] 键 raft::device_vector_view (len = ncols)。假定此数组中的每个键都在 [0, nkeys) 之间。如果不是这样，调用者应已调用 make_monotonic 基本操作来准备此类连续且单调递增的键数组。

  • out – [out] 沿列归约后的输出 raft::device_matrix_view (dim = nrows x nkeys)。假定其为行主布局

  • nkeys – [in] keys 数组中唯一键的数量。默认情况下，从 out 的列数推断

  • reset_sums – [in] 在归约之前是否将输出总和重置为零

按键归约行#

#include <raft/linalg/reduce_rows_by_key.cuh>

namespace raft::linalg

template<typename ElementType, typename KeyType, typename WeightType, typename IndexType>
void reduce_rows_by_key(raft::resources const &handle, raft::device_matrix_view<const ElementType, IndexType, raft::row_major> d_A, raft::device_vector_view<const KeyType, IndexType> d_keys, raft::device_matrix_view<ElementType, IndexType, raft::row_major> d_sums, IndexType n_unique_keys, raft::device_vector_view<char, IndexType> d_keys_char, std::optional<raft::device_vector_view<const WeightType, IndexType>> d_weights = std::nullopt, bool reset_sums = true)#

计算给定键的矩阵行的加权求和归约。TODO: 支持通用归约 lambda rapidsai/raft#860。

模板参数:

  • ElementType – 输入和输出的数据类型

  • KeyType – 键的数据类型

  • WeightType – 权重的数据类型

  • IndexType – 索引类型

参数:

带步长的归约#

#include <raft/linalg/strided_reduction.cuh>

namespace raft::linalg

template<typename InValueType, typename LayoutPolicy, typename OutValueType, typename IndexType, typename MainLambda = raft::identity_op, typename ReduceLambda = raft::add_op, typename FinalLambda = raft::identity_op>
void strided_reduction(raft::resources const &handle, raft::device_matrix_view<const InValueType, IndexType, LayoutPolicy> data, raft::device_vector_view<OutValueType, IndexType> dots, OutValueType init, bool inplace = false, MainLambda main_op = raft::identity_op(), ReduceLambda reduce_op = raft::add_op(), FinalLambda final_op = raft::identity_op())#

沿带步长的维度计算输入矩阵的归约。当期望的归约不是沿内存布局的维度时，可以使用此 API。例如，行主矩阵将沿行归约，而列主矩阵将沿列归约。

模板参数:

  • InValueType – 底层 raft::matrix_view 的输入数据类型

  • LayoutPolicy – 输入/输出的布局（行主序或列主序）

  • OutValueType – 底层 raft::matrix_view 和归约的输出数据类型

  • IndexType – 用于寻址的整数类型

  • MainLambda – 累加时应用的一元 lambda 函数（例如：L1 或 L2 范数）。它必须是一个支持以下输入和输出的“可调用对象”。

  • ReduceLambda – 用于归约的二元 lambda 函数（例如：L2 范数的加法（+））。它必须是一个支持以下输入和输出的“可调用对象”。

  • FinalLambda – 在 STG 前应用的最终 lambda 函数（例如：L2 范数的开方）。它必须是一个支持以下输入和输出的“可调用对象”。

参数:

  • handle – [in] raft::resources

  • data – [in] 类型为 raft::device_matrix_view 的输入

  • dots – [out] 类型为 raft::device_matrix_view 的输出

  • init – [in] 用于归约的初始值

  • main_op – [in] 在归约前应用的融合逐元素运算

  • reduce_op – [in] 融合的二元归约运算

  • final_op – [in] 在存储结果前应用的融合逐元素运算

  • inplace – [in] 归约结果是在原位添加还是覆盖旧值？