基于CUDA从含重复三元组的设备指针组装COO/CSR稀疏矩阵

先明确：cuSparse的COO格式不原生支持重复的(i,j)条目，但可以通过「合并重复项 + 格式转换/创建」的流程解决你的需求，具体步骤如下：

1. 合并重复的(i,j)条目 #

你的设备端三元组存在重复索引，第一步必须对相同(i,j)位置的值进行累加，这里推荐两种实现方式：

• 用Thrust库简化操作：先通过thrust::sort_by_key按照(row, col)键对所有三元组排序，再用thrust::reduce_by_key对相同键对应的值做累加，直接得到去重后的COO数组。这种方式无需手写复杂内核，适合大多数场景。
• 自定义CUDA内核：如果对性能有极致要求，可以手写排序+合并内核，不过Thrust的实现已经经过优化，除非有特殊需求，不建议重复造轮子。

2. 构建规范的COO或CSR矩阵 #

得到去重且行索引非降序的COO数组后，就可以用cuSparse的原生函数处理了：

• 创建COO矩阵：调用cusparseCreateCoo，传入去重后的行、列、值指针，以及矩阵的行列数、非零元素数即可。注意cuSparse要求COO的行索引必须是非降序排列的，这一步在之前的排序环节已经满足。
• 直接创建CSR矩阵：调用cusparseXcoo2csr将去重排序后的COO格式转换为CSR格式，这是cuSparse的高效原生操作，适合后续需要CSR格式的计算场景。

核心代码示例 #

// 假设已有设备端指针：d_rows(行索引), d_cols(列索引), d_vals(值)，原始非零元素数nnz_raw
// 矩阵行数m，列数n

// 1. 用Thrust按(row, col)排序
thrust::device_ptr<int> dev_rows(d_rows);
thrust::device_ptr<int> dev_cols(d_cols);
thrust::device_ptr<double> dev_vals(d_vals);

// 以(row, col)为键排序值数组
thrust::sort_by_key(
    thrust::make_zip_iterator(thrust::make_tuple(dev_rows, dev_cols)),
    thrust::make_zip_iterator(thrust::make_tuple(dev_rows + nnz_raw, dev_cols + nnz_raw)),
    dev_vals
);

// 2. 合并重复项，得到去重后的COO
// 预先分配去重后的设备内存（可以先估算或用Thrust的临时容器）
thrust::device_vector<int> new_rows(nnz_raw);
thrust::device_vector<int> new_cols(nnz_raw);
thrust::device_vector<double> new_vals(nnz_raw);

auto [key_end, val_end] = thrust::reduce_by_key(
    thrust::make_zip_iterator(thrust::make_tuple(dev_rows, dev_cols)),
    thrust::make_zip_iterator(thrust::make_tuple(dev_rows + nnz_raw, dev_cols + nnz_raw)),
    dev_vals,
    thrust::make_zip_iterator(thrust::make_tuple(new_rows.begin(), new_cols.begin())),
    new_vals.begin(),
    // 判断两个键是否相等（即(row,col)是否相同）
    [] __host__ __device__ (const auto& a, const auto& b) {
        return a == b;
    },
    // 累加相同键的值
    thrust::plus<double>()
);
int nnz_unique = key_end - thrust::make_zip_iterator(thrust::make_tuple(new_rows.begin(), new_cols.begin()));

// 3. 转换为CSR或创建COO
cusparseHandle_t cusparse_h;
cusparseCreate(&cusparse_h);
cusparseMatDescr_t mat_descr;
cusparseCreateMatDescr(&mat_descr);
cusparseSetMatIndexBase(mat_descr, CUSPARSE_INDEX_BASE_ZERO);

// 方式一：创建CSR矩阵
int* d_csr_row_ptr, *d_csr_col_idx;
double* d_csr_vals;
// 分配CSR行指针内存（大小为m+1）
cudaMalloc(&d_csr_row_ptr, (m + 1) * sizeof(int));
cudaMalloc(&d_csr_col_idx, nnz_unique * sizeof(int));
cudaMalloc(&d_csr_vals, nnz_unique * sizeof(double));

// COO转CSR
cusparseXcoo2csr(
    cusparse_h,
    thrust::raw_pointer_cast(new_rows.data()),
    nnz_unique,
    m,
    d_csr_row_ptr,
    CUSPARSE_INDEX_BASE_ZERO
);
// 复制列索引和值到CSR
cudaMemcpy(d_csr_col_idx, thrust::raw_pointer_cast(new_cols.data()), nnz_unique * sizeof(int), cudaMemcpyDeviceToDevice);
cudaMemcpy(d_csr_vals, thrust::raw_pointer_cast(new_vals.data()), nnz_unique * sizeof(double), cudaMemcpyDeviceToDevice);

// 方式二：创建COO矩阵
cusparseSpMatDescr_t coo_mat;
cusparseCreateCoo(
    &coo_mat,
    m, n, nnz_unique,
    thrust::raw_pointer_cast(new_rows.data()),
    thrust::raw_pointer_cast(new_cols.data()),
    thrust::raw_pointer_cast(new_vals.data()),
    CUSPARSE_INDEX_32I, CUSPARSE_INDEX_32I,
    CUSPARSE_INDEX_BASE_ZERO, CUDA_R_64F
);

// 后续操作...
// 记得释放资源
cusparseDestroySpMat(coo_mat);
cusparseDestroyMatDescr(mat_descr);
cusparseDestroy(cusparse_h);
cudaFree(d_csr_row_ptr);
cudaFree(d_csr_col_idx);
cudaFree(d_csr_vals);

关键注意点 #

• 排序是硬性要求：cuSparse的COO和CSR转换都要求行索引非降序，否则会返回错误。
• 内存分配：去重后的非零元素数nnz_unique会小于等于原始的nnz_raw，可以用Thrust的返回值直接获取，避免内存浪费。
• 索引基：确保代码中统一使用0基或1基（示例中用0基，和Eigen默认一致），避免索引越界。

内容的提问来源于stack exchange，提问作者jjcasmar