注意

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

稀疏线性代数

template<typename T>
size_t raft::sparse::linalg::csr_add_calc_inds(const int *a_ind, const int *a_indptr, const T *a_val, int nnz1, const int *b_ind, const int *b_indptr, const T *b_val, int nnz2, int m, int *out_ind, cudaStream_t stream)

计算将两个 CSR 矩阵相加后得到的 CSR row_ind 数组。

参数：

• a_ind – 左侧 row_ind 数组

• a_indptr – 左侧 index_ptr 数组

• a_val – 左侧数据数组

• nnz1 – 左侧 index_ptr 和 val 数组的大小

• b_ind – 右侧 row_ind 数组

• b_indptr – 右侧 index_ptr 数组

• b_val – 右侧数据数组

• nnz2 – 右侧 index_ptr 和 val 数组的大小

• m – 输出数组的大小（最终矩阵的行数）

• out_ind – 输出 row_ind 数组

• stream – 要使用的 cuda stream

template<typename T>
void raft::sparse::linalg::csr_add_finalize(const int *a_ind, const int *a_indptr, const T *a_val, int nnz1, const int *b_ind, const int *b_indptr, const T *b_val, int nnz2, int m, int *c_ind, int *c_indptr, T *c_val, cudaStream_t stream)

计算将两个 CSR 矩阵相加后得到的 CSR row_ind 数组。

参数：

• a_ind – 左侧 row_ind 数组

• a_indptr – 左侧 index_ptr 数组

• a_val – 左侧数据数组

• nnz1 – 左侧 index_ptr 和 val 数组的大小

• b_ind – 右侧 row_ind 数组

• b_indptr – 右侧 index_ptr 数组

• b_val – 右侧数据数组

• nnz2 – 右侧 index_ptr 和 val 数组的大小

• m – 输出数组的大小（最终矩阵的行数）

• c_ind – 输出 row_ind 数组

• c_indptr – 输出 ind_ptr 数组

• c_val – 输出数据数组

• stream – 要使用的 cuda stream

template<typename T = int, typename nnz_type, typename outT>
void raft::sparse::linalg::coo_degree(const T *rows, nnz_type nnz, outT *results, cudaStream_t stream)

统计每一行的值的数量。

模板参数：

TPB_X – 每个块使用的线程数

参数：

• rows – COO 矩阵的 rows 数组

• nnz – rows 数组的大小

• results – 输出结果数组

• stream – 要使用的 cuda stream

template<typename T, typename outT>
void raft::sparse::linalg::coo_degree(COO<T> *in, outT *results, cudaStream_t stream)

统计每一行的值的数量。

模板参数：

• TPB_X – 每个块使用的线程数

• T – 底层 values 数组的类型名称

参数：

• in – 用于统计行数的输入 COO 对象

• results – 包含行数的输出数组（大小=in->n_rows）

• stream – 要使用的 cuda stream

template<typename T, typename nnz_type, typename outT>
void raft::sparse::linalg::coo_degree_scalar(const int *rows, const T *vals, nnz_type nnz, T scalar, outT *results, cudaStream_t stream = 0)

统计每一行中与特定标量不匹配的值的数量。

模板参数：

• TPB_X – 每个块使用的线程数

• T – 底层 value 数组的类型名称

参数：

• rows – 输入 COO row 数组

• vals – 输入 COO val 数组

• nnz – 输入 COO 数组的大小

• scalar – 用于统计行数的匹配标量

• results – 输出行数

• stream – 要使用的 cuda stream

template<typename T, typename outT>
void raft::sparse::linalg::coo_degree_scalar(COO<T> *in, T scalar, outT *results, cudaStream_t stream)

统计每一行中与特定标量不匹配的值的数量。

模板参数：

• TPB_X – 每个块使用的线程数

• T – 底层 value 数组的类型名称

参数：

• in – 输入 COO 数组

• scalar – 用于统计行数的匹配标量

• results – 输出行数

• stream – 要使用的 cuda stream

template<typename T>
void raft::sparse::linalg::coo_degree_nz(const int *rows, const T *vals, int nnz, int *results, cudaStream_t stream)

统计每一行的非零元素数量。

模板参数：

• TPB_X – 每个块使用的线程数

• T – 底层 value 数组的类型名称

参数：

• rows – 输入 COO row 数组

• vals – 输入 COO val 数组

• nnz – 输入 COO 数组的大小

• results – 输出行数

• stream – 要使用的 cuda stream

template<typename T>
void raft::sparse::linalg::coo_degree_nz(COO<T> *in, int *results, cudaStream_t stream)

统计每一行的非零值的数量。

模板参数：

• TPB_X – 每个块使用的线程数

• T – 底层 value 数组的类型名称

参数：

• in – 输入 COO 数组

• results – 输出行数

• stream – 要使用的 cuda stream

template<typename ElementType, typename IndptrType, typename IndicesType, typename NZType>
auto raft::sparse::linalg::compute_graph_laplacian(raft::resources const &res, device_csr_matrix_view<ElementType, IndptrType, IndicesType, NZType> input)

给定一个 CSR 邻接矩阵，返回图 Laplacian

注意，对于非对称矩阵，返回的是出度 Laplacian。

template<typename T, typename indT>
void raft::sparse::linalg::csr_row_normalize_l1(const indT *ia, const T *vals, indT nnz, int m, T *result, cudaStream_t stream)

对给定 CSR 格式的稀疏矩阵的行执行 L1 范数归一化。

参数：

• ia – row_ind 数组

• vals – 数据数组

• nnz – 数据数组的大小

• m – row_ind 数组的大小

• result – L1 范数归一化后的数据数组

• stream – 要使用的 cuda stream

template<typename T>
void raft::sparse::linalg::csr_row_normalize_max(const int *ia, const T *vals, int nnz, int m, T *result, cudaStream_t stream)

对给定 CSR 格式的稀疏矩阵执行 L_inf 范数归一化。

参数：

• ia – row_ind 数组

• vals – 数据数组

• nnz – 数据数组的大小

• m – row_ind 数组的大小

• result – L1 范数归一化后的数据数组

• stream – 要使用的 cuda stream

template<typename Type, typename IdxType = int, typename Lambda = raft::identity_op>
void raft::sparse::linalg::rowNormCsr(raft::resources const &handle, const IdxType *ia, const Type *data, const IdxType nnz, const IdxType N, Type *norm, raft::linalg::NormType type, Lambda fin_op = raft::identity_op())

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

行范数在计算成对距离矩阵时很有用，例如在许多聚类算法（如 knn, kmeans, dbscan 等）中使用。

模板参数：

• Type – 数据类型

• Lambda – device final lambda

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

参数：

• handle – [in] raft 句柄

• ia – 输入矩阵的行索引数组

• data – 输入矩阵的非零元素数据

• nnz – 数据中的元素数量

• N – 行数

• norm – 行范数的输出向量，大小为 [N]

• type – 要应用的范数类型

• fin_op – 最终 lambda 操作

template<typename ValueType, typename IndexType, typename NZType, typename LayoutPolicyA, typename LayoutPolicyB, typename OutputType>
void raft::sparse::linalg::sddmm(raft::resources const &handle, raft::device_matrix_view<const ValueType, IndexType, LayoutPolicyA> A, raft::device_matrix_view<const ValueType, IndexType, LayoutPolicyB> B, raft::device_csr_matrix_view<OutputType, IndexType, IndexType, NZType> C, const raft::linalg::Operation opA, const raft::linalg::Operation opB, raft::host_scalar_view<OutputType> alpha, raft::host_scalar_view<OutputType> beta)

此函数执行密集矩阵 A 和密集矩阵 B 的乘法，然后与 C 的稀疏模式进行元素乘法。它计算以下等式：C = alpha · (opA(A) * opB(B) ∘ spy(C)) + beta · C 其中 A、B 是设备矩阵视图，C 是 CSR 设备矩阵视图。

模板参数：

• ValueType – 输入/输出矩阵的数据类型 (float/double/half)

• IndexType – C 的类型

• NZType – C 的类型

• LayoutPolicyA – A 的布局

• LayoutPolicyB – B 的布局

• OutputType – 输出类型，默认为 ValueType

参数：

• handle – [in] raft 句柄

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

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

• C – [inout] 输出 raft::device_csr_matrix_view

• opA – [in] 输入操作 op(A)

• opB – [in] 输入操作 op(B)

• alpha – [in] 输入 raft::host_scalar_view

• beta – [in] 输入 raft::host_scalar_view

template<typename ValueType, typename IndexType, typename NZType, typename LayoutPolicyY, typename LayoutPolicyZ>
void raft::sparse::linalg::spmm(raft::resources const &handle, const bool trans_x, const bool trans_y, const ValueType *alpha, raft::device_csr_matrix_view<const ValueType, int, int, NZType> x, raft::device_matrix_view<const ValueType, IndexType, LayoutPolicyY> y, const ValueType *beta, raft::device_matrix_view<ValueType, IndexType, LayoutPolicyZ> z)

SPMM 函数旨在处理 cuSparse 中所有 CSR * DENSE 操作数布局组合。它计算以下等式：Z = alpha . X * Y + beta . Z 其中 X 是 CSR 设备矩阵视图，Y,Z 是设备矩阵视图。

模板参数：

• ValueType – 输入/输出矩阵的数据类型 (float/double)

• IndexType – Y 和 Z 的类型

• NZType – X 的类型

• LayoutPolicyY – Y 的布局

• LayoutPolicyZ – Z 的布局

参数：

• handle – [in] raft 句柄

• trans_x – [in] X 的转置操作

• trans_y – [in] Y 的转置操作

• alpha – [in] 标量

• x – [in] 输入 raft::device_csr_matrix_view

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

• beta – [in] 标量

• z – [inout] 输入-输出 raft::device_matrix_view

template<typename T, typename Lambda>
void raft::sparse::linalg::coo_symmetrize(COO<T> *in, COO<T> *out, Lambda reduction_op, cudaStream_t stream)

接受一个可能不对称的 COO 矩阵并对其进行对称化，对每个值及其转置值运行自定义归约函数。

参数：

• in – 输入 COO 矩阵

• out – 输出对称化后的 COO 矩阵

• reduction_op – 自定义归约函数

• stream – 要使用的 cuda stream

template<typename value_idx = int64_t, typename value_t = float> RAFT_KERNEL raft::sparse::linalg::symmetric_find_size (const value_t __restrict__ *data, const value_idx __restrict__ *indices, const value_idx n, const int k, value_idx __restrict__ *row_sizes, value_idx __restrict__ *row_sizes2)

查找每行所需的空间量。我们遍历所有数据点并增加每行的计数。

TODO: 这不是通用的。替换 symmetrize() 来移除它

参数：

• data – 输入 knn 距离(n, k)

• indices – 输入 knn 索引(n, k)

• n – 行数

• k – 邻居数

• row_sizes – 输入空行求和数组 1 (n)

• row_sizes2 – 输入空行求和数组 2 (n)，用于更快的归约

template<typename value_idx> RAFT_KERNEL raft::sparse::linalg::reduce_find_size (const value_idx n, const int k, value_idx __restrict__ *row_sizes, const value_idx __restrict__ *row_sizes2)

对 symmetric_find_size 内核执行 sum(row_sizes) + k 归约。使算法更快。

TODO: 这不是通用的。替换 symmetrize() 来移除它

参数：

• n – 行数

• k – 邻居数

• row_sizes – 输入行求和数组 1 (n)

• row_sizes2 – 输入行求和数组 2 (n)，用于更快的归约

template<typename value_idx = int64_t, typename value_t = float> RAFT_KERNEL raft::sparse::linalg::symmetric_sum (value_idx *__restrict__ edges, const value_t *__restrict__ data, const value_idx __restrict__ *indices, value_t *__restrict__ VAL, value_idx *__restrict__ COL, value_idx *__restrict__ ROW, const value_idx n, const int k)

执行 data + data.T 操作。只有在确定 data + data.T 的 CSR 矩阵的 row_sizes 后才能运行。

TODO: 这不是通用的。替换 symmetrize() 来移除它

参数：

• edges – 归约后的输入行求和数组 (n)

• data – 输入 knn 距离(n, k)

• indices – 输入 knn 索引(n, k)

• VAL – data + data.T 的输出值

• COL – data + data.T 的输出列索引

• ROW – data + data.T 的输出行索引

• n – 行数

• k – 邻居数

template<typename value_idx = int64_t, typename value_t = float, int TPB_X = 32, int TPB_Y = 32> void raft::sparse::linalg::from_knn_symmetrize_matrix (const value_idx __restrict__ knn_indices, const value_t __restrict__ knn_dists, const value_idx n, const int k, COO< value_t, value_idx > *out, cudaStream_t stream)

对原始 KNN 数据执行 data + data.T 操作。调用以下步骤：(1) 查找每行所需的空间量 (2) 计算所需的最终空间 (n*k + sum(row_sizes)) == 2*n*k (3) 分配新空间 (4) 为每个新行准备边 (5) 执行最终的 data + data.T 操作 (6) 返回求和后的 VAL, COL, ROW。

TODO: 这不是通用的。替换 symmetrize() 来移除它

参数：

• knn_indices – 输入 knn 距离(n, k)

• knn_dists – 输入 knn 索引(n, k)

• n – 行数

• k – 邻居数

• out – 输出 COO Matrix 类

• stream – 输入 cuda stream

template<typename value_idx, typename value_t, typename nnz_t>
void raft::sparse::linalg::symmetrize(raft::resources const &handle, const value_idx *rows, const value_idx *cols, const value_t *vals, value_idx m, value_idx n, nnz_t nnz, raft::sparse::COO<value_t, value_idx, nnz_t> &out)

将COO矩阵对称化

template<typename value_idx, typename value_t>
void raft::sparse::linalg::csr_transpose(raft::resources const &handle, const value_idx *csr_indptr, const value_idx *csr_indices, const value_t *csr_data, value_idx *csc_indptr, value_idx *csc_indices, value_t *csc_data, value_idx csr_nrows, value_idx csr_ncols, value_idx nnz, cudaStream_t stream)

将一组 CSR 数组转置为一组 CSC 数组。

模板参数：

• value_idx – : CSR 索引数组的数据类型

• value_t – : CSR 数据数组的数据类型

参数：

• handle – [in] : 用于调用 cusparse

• csr_indptr – [in] : CSR 行索引数组

• csr_indices – [in] : CSR 列索引数组

• csr_data – [in] : CSR 数据数组

• csc_indptr – [out] : CSC 行索引数组

• csc_indices – [out] : CSC 列索引数组

• csc_data – [out] : CSC 数据数组

• csr_nrows – [in] : CSR 中的行数

• csr_ncols – [in] : CSR 中的列数

• nnz – [in] : CSR 中的非零元素数量

• stream – [in] : 用于排序事件的 Cuda stream