CANN/ops-sparse SpMM算子

SpMM算子

【免费下载链接】ops-sparse本项目是CANN提供的高性能稀疏矩阵计算的算子库，专注于优化稀疏矩阵的计算效率。项目地址: https://gitcode.com/cann/ops-sparse

算子概述

SpMM（Sparse Matrix - Dense Matrix Multiplication）算子实现稀疏矩阵与稠密矩阵的乘法运算。核心运算为 C = alpha * op(A) * op(B) + beta * C，其中 A 为 CSR 格式稀疏矩阵，B 和 C 为稠密矩阵。

数学表达式：

C = alpha * op(A) * op(B) + beta * C

其中 op(A) 为稀疏矩阵 A 的操作（转置/非转置），op(B) 为稠密矩阵 B 的操作。

调用流程为三步法：

1. GetBufferSize：查询所需 workspace 大小
2. Preprocess：对稀疏矩阵进行预处理（行重排 + 分箱），加速后续计算
3. SpMM：执行 C = alpha * op(A) * op(B) + beta * C

包含以下接口：

算子执行接口

aclsparseSpMMGetBufferSize

产品支持情况

• Ascend 950PR / Ascend 950DT：支持
• Atlas A3 训练系列产品 / Atlas A3 推理系列产品：不支持
• Atlas A2 训练系列产品 / Atlas A2 推理系列产品：不支持

函数原型

aclsparseStatus_t aclsparseSpMMGetBufferSize(aclsparseHandle_t handle, aclsparseOperation_t opA, aclsparseOperation_t opB, const void *alpha, aclsparseConstSpMatDescr_t matA, aclsparseConstDnMatDescr_t matB, const void *beta, aclsparseDnMatDescr_t matC, aclDataType computeType, aclsparseSpMMAlg_t alg, size_t *size)

参数说明

约束说明

• handle 不可为 nullptr
• matA、matB、matC 不可为 nullptr
• matA 仅支持 CSR 格式（ACL_SPARSE_FORMAT_CSR）
• matA 的行偏移和列索引类型必须均为ACL_SPARSE_INDEX_32I，且两者类型相同
• matA 的索引基址仅支持ACL_SPARSE_INDEX_BASE_ZERO
• opA 仅支持ACL_SPARSE_OP_NON_TRANSPOSE
• opB 支持ACL_SPARSE_OP_NON_TRANSPOSE或ACL_SPARSE_OP_TRANSPOSE，不支持共轭转置
• 数据类型组合仅支持以下三种：
  • matA=ACL_FLOAT, matB=ACL_FLOAT, matC=ACL_FLOAT, computeType=ACL_FLOAT
  • matA=ACL_FLOAT16, matB=ACL_FLOAT16, matC=ACL_FLOAT16, computeType=ACL_FLOAT
  • matA=ACL_INT8, matB=ACL_INT8, matC=ACL_INT32, computeType=ACL_INT32
• 维度匹配：A.cols == B.rows，A.rows == C.rows，B.cols == C.cols
• size 不可为 nullptr

支持的稀疏格式

aclsparseSpMMPreprocess

产品支持情况

• Ascend 950PR / Ascend 950DT：支持
• Atlas A3 训练系列产品 / Atlas A3 推理系列产品：不支持
• Atlas A2 训练系列产品 / Atlas A2 推理系列产品：不支持

函数原型

aclsparseStatus_t aclsparseSpMMPreprocess(aclsparseHandle_t handle, aclsparseOperation_t opA, aclsparseOperation_t opB, const void *alpha, aclsparseConstSpMatDescr_t matA, aclsparseConstDnMatDescr_t matB, const void *beta, aclsparseDnMatDescr_t matC, aclDataType computeType, aclsparseSpMMAlg_t alg, void *buffer)

参数说明

约束说明

• 同 aclsparseSpMMGetBufferSize 的约束
• buffer 不可为 nullptr，需按 GetBufferSize 返回的大小分配

支持的稀疏格式

aclsparseSpMM

产品支持情况

• Ascend 950PR / Ascend 950DT：支持
• Atlas A3 训练系列产品 / Atlas A3 推理系列产品：不支持
• Atlas A2 训练系列产品 / Atlas A2 推理系列产品：不支持

函数原型

aclsparseStatus_t aclsparseSpMM(aclsparseHandle_t handle, aclsparseOperation_t opA, aclsparseOperation_t opB, const void *alpha, aclsparseConstSpMatDescr_t matA, aclsparseConstDnMatDescr_t matB, const void *beta, aclsparseDnMatDescr_t matC, aclDataType computeType, aclsparseSpMMAlg_t alg, void *buffer)

参数说明

约束说明

• 同 aclsparseSpMMGetBufferSize 的约束
• buffer 不可为 nullptr，需按 GetBufferSize 返回的大小分配
• 调用前须先调用aclsparseSpMMPreprocess进行预处理
• ACL_SPARSE_SPMM_CSR_FP32_HIGH_PRECISION_ALG仅在 computeType=ACL_FLOAT 时生效（使用 Kahan 补偿求和提升精度），对 fp16/int8 静默忽略

支持的稀疏格式

调用示例

示例代码如下，仅供参考，具体编译和执行过程请参考编译与运行样例。

#include <cstdio> #include <memory> #include <vector> #include "acl/acl.h" #include "cann_ops_sparse.h" #define CHECK_RET(cond, return_expr) \ do { \ if (!(cond)) { \ return_expr; \ } \ } while (0) #define LOG_PRINT(message, ...) \ do { \ printf(message, ##__VA_ARGS__); \ } while (0) class AclContext { public: explicit AclContext(int32_t deviceId) : deviceId_(deviceId) {} ~AclContext() { if (stream_ != nullptr) { aclrtDestroyStream(stream_); stream_ = nullptr; } if (deviceSet_) { aclrtResetDevice(deviceId_); deviceSet_ = false; } if (aclInited_) { aclFinalize(); aclInited_ = false; } } int Init() { auto ret = aclInit(nullptr); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclInit failed. ERROR: %d\n", ret); return ret); aclInited_ = true; ret = aclrtSetDevice(deviceId_); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); return ret); deviceSet_ = true; ret = aclrtCreateStream(&stream_); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateStream failed. ERROR: %d\n", ret); return ret); return ACL_SUCCESS; } aclrtStream Stream() const { return stream_; } private: int32_t deviceId_; aclrtStream stream_ = nullptr; bool aclInited_ = false; bool deviceSet_ = false; }; // 辅助：分配 Device 内存并拷贝 Host 数据 static void* AllocAndCopyDevice(const void *hostPtr, size_t sizeBytes) { void *dPtr = nullptr; aclrtMalloc(&dPtr, sizeBytes, ACL_MEM_MALLOC_HUGE_FIRST); if (hostPtr != nullptr && sizeBytes > 0) { aclrtMemcpy(dPtr, sizeBytes, hostPtr, sizeBytes, ACL_MEMCPY_HOST_TO_DEVICE); } return dPtr; } int aclsparseSpMMTest(AclContext& ctx) { aclrtStream stream = ctx.Stream(); // 1. 创建 ops-sparse 句柄 aclsparseHandle_t rawHandle = nullptr; auto sparseRet = aclsparseCreate(&rawHandle); CHECK_RET(sparseRet == ACL_SPARSE_STATUS_SUCCESS, LOG_PRINT("aclsparseCreate failed. ERROR: %d\n", sparseRet); return sparseRet); std::unique_ptr<aclsparseContext, aclsparseStatus_t (*)(aclsparseHandle_t)> handlePtr(rawHandle, aclsparseDestroy); sparseRet = aclsparseSetStream(static_cast<aclsparseHandle_t>(handlePtr.get()), stream); CHECK_RET(sparseRet == ACL_SPARSE_STATUS_SUCCESS, LOG_PRINT("aclsparseSetStream failed. ERROR: %d\n", sparseRet); return sparseRet); // 2. 设置 PointerMode sparseRet = aclsparseSetPointerMode(static_cast<aclsparseHandle_t>(handlePtr.get()), ACL_SPARSE_POINTER_MODE_HOST); CHECK_RET(sparseRet == ACL_SPARSE_STATUS_SUCCESS, LOG_PRINT("aclsparseSetPointerMode failed. ERROR: %d\n", sparseRet); return sparseRet); // 3. 准备 Host 端 CSR 数据 // A (4x3, nnz=5): B (3x2): // [1.0 0.0 2.0] [1.0 0.0] // [0.0 3.0 0.0] [0.0 1.0] // [4.0 0.0 0.0] [1.0 1.0] // [0.0 0.0 5.0] // // C = 1.0 * A * B (4x2): // C[0] = [3.0, 2.0] // C[1] = [0.0, 3.0] // C[2] = [4.0, 0.0] // C[3] = [5.0, 5.0] int64_t m = 4, k = 3, n = 2; int64_t nnzA = 5; float hAlpha = 1.0f; float hBeta = 0.0f; std::vector<int> hRowPtrA = {0, 2, 3, 4, 5}; std::vector<int> hColIndA = {0, 2, 1, 0, 2}; std::vector<float> hValA = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f}; // B: 行主序 3x2 int64_t ldb = n, ldc = n; aclsparseOrder_t orderB = ACL_SPARSE_ORDER_ROW; aclsparseOrder_t orderC = ACL_SPARSE_ORDER_ROW; std::vector<float> hB(static_cast<size_t>(k) * n, 0.0f); hB[0 * n + 0] = 1.0f; hB[0 * n + 1] = 0.0f; hB[1 * n + 0] = 0.0f; hB[1 * n + 1] = 1.0f; hB[2 * n + 0] = 1.0f; hB[2 * n + 1] = 1.0f; std::vector<float> hC(static_cast<size_t>(m) * n, 0.0f); // 4. 拷贝数据到 Device void *dRowPtrA = AllocAndCopyDevice(hRowPtrA.data(), (m + 1) * sizeof(int)); void *dColIndA = AllocAndCopyDevice(hColIndA.data(), nnzA * sizeof(int)); void *dValA = AllocAndCopyDevice(hValA.data(), nnzA * sizeof(float)); void *dB = AllocAndCopyDevice(hB.data(), static_cast<size_t>(k) * n * sizeof(float)); void *dC = AllocAndCopyDevice(hC.data(), static_cast<size_t>(m) * n * sizeof(float)); // 5. 创建描述符 aclsparseSpMatDescr_t matA = nullptr; sparseRet = aclsparseCreateCsr(&matA, m, k, nnzA, dRowPtrA, dColIndA, dValA, ACL_SPARSE_INDEX_32I, ACL_SPARSE_INDEX_32I, ACL_SPARSE_INDEX_BASE_ZERO, ACL_FLOAT); CHECK_RET(sparseRet == ACL_SPARSE_STATUS_SUCCESS, LOG_PRINT("aclsparseCreateCsr failed. ERROR: %d\n", sparseRet); return sparseRet); aclsparseDnMatDescr_t matB = nullptr, matC = nullptr; sparseRet = aclsparseCreateDnMat(&matB, k, n, ldb, dB, ACL_FLOAT, orderB); CHECK_RET(sparseRet == ACL_SPARSE_STATUS_SUCCESS, LOG_PRINT("aclsparseCreateDnMat B failed. ERROR: %d\n", sparseRet); return sparseRet); sparseRet = aclsparseCreateDnMat(&matC, m, n, ldc, dC, ACL_FLOAT, orderC); CHECK_RET(sparseRet == ACL_SPARSE_STATUS_SUCCESS, LOG_PRINT("aclsparseCreateDnMat C failed. ERROR: %d\n", sparseRet); return sparseRet); // 6. Step 1 — GetBufferSize size_t bufferSize = 0; sparseRet = aclsparseSpMMGetBufferSize( static_cast<aclsparseHandle_t>(handlePtr.get()), ACL_SPARSE_OP_NON_TRANSPOSE, ACL_SPARSE_OP_NON_TRANSPOSE, &hAlpha, matA, matB, &hBeta, matC, ACL_FLOAT, ACL_SPARSE_SPMM_CSR_ALG1, &bufferSize); CHECK_RET(sparseRet == ACL_SPARSE_STATUS_SUCCESS, LOG_PRINT("SpMMGetBufferSize failed. ERROR: %d\n", sparseRet); return sparseRet); LOG_PRINT("bufferSize = %zu bytes\n", bufferSize); void *dBuffer = nullptr; auto aclRet = aclrtMalloc(&dBuffer, bufferSize, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(aclRet == ACL_SUCCESS, LOG_PRINT("aclrtMalloc for buffer failed. ERROR: %d\n", aclRet); return aclRet); // 7. Step 2 — Preprocess sparseRet = aclsparseSpMMPreprocess( static_cast<aclsparseHandle_t>(handlePtr.get()), ACL_SPARSE_OP_NON_TRANSPOSE, ACL_SPARSE_OP_NON_TRANSPOSE, &hAlpha, matA, matB, &hBeta, matC, ACL_FLOAT, ACL_SPARSE_SPMM_CSR_ALG1, dBuffer); CHECK_RET(sparseRet == ACL_SPARSE_STATUS_SUCCESS, LOG_PRINT("SpMMPreprocess failed. ERROR: %d\n", sparseRet); return sparseRet); // 8. Step 3 — SpMM sparseRet = aclsparseSpMM( static_cast<aclsparseHandle_t>(handlePtr.get()), ACL_SPARSE_OP_NON_TRANSPOSE, ACL_SPARSE_OP_NON_TRANSPOSE, &hAlpha, matA, matB, &hBeta, matC, ACL_FLOAT, ACL_SPARSE_SPMM_CSR_ALG1, dBuffer); CHECK_RET(sparseRet == ACL_SPARSE_STATUS_SUCCESS, LOG_PRINT("SpMM failed. ERROR: %d\n", sparseRet); return sparseRet); // 9. 同步等待计算完成 aclRet = aclrtSynchronizeStream(stream); CHECK_RET(aclRet == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", aclRet); return aclRet); // 10. 将结果拷贝回 Host 并打印 aclRet = aclrtMemcpy(hC.data(), static_cast<size_t>(m) * n * sizeof(float), dC, static_cast<size_t>(m) * n * sizeof(float), ACL_MEMCPY_DEVICE_TO_HOST); CHECK_RET(aclRet == ACL_SUCCESS, LOG_PRINT("copy result from device to host failed. ERROR: %d\n", aclRet); return aclRet); for (int64_t i = 0; i < m; i++) { LOG_PRINT("C[%lld] = %.1f, %.1f\n", static_cast<long long>(i), hC[i * n + 0], hC[i * n + 1]); } // 11. 清理资源 aclsparseDestroySpMat(matA); aclsparseDestroyDnMat(matB); aclsparseDestroyDnMat(matC); if (dRowPtrA) aclrtFree(dRowPtrA); if (dColIndA) aclrtFree(dColIndA); if (dValA) aclrtFree(dValA); if (dB) aclrtFree(dB); if (dC) aclrtFree(dC); if (dBuffer) aclrtFree(dBuffer); return ACL_SPARSE_STATUS_SUCCESS; } int main() { AclContext ctx(0); auto ret = ctx.Init(); CHECK_RET(ret == ACL_SUCCESS, return ret); ret = aclsparseSpMMTest(ctx); CHECK_RET(ret == ACL_SPARSE_STATUS_SUCCESS, LOG_PRINT("aclsparseSpMMTest failed. ERROR: %d\n", ret); return ret); return 0; }

预期输出如下：

bufferSize = 512 bytes C[0] = 3.0, 2.0 C[1] = 0.0, 3.0 C[2] = 4.0, 0.0 C[3] = 5.0, 5.0

【免费下载链接】ops-sparse本项目是CANN提供的高性能稀疏矩阵计算的算子库，专注于优化稀疏矩阵的计算效率。项目地址: https://gitcode.com/cann/ops-sparse