SpMM算子【免费下载链接】ops-sparse本项目是CANN提供的高性能稀疏矩阵计算的算子库专注于优化稀疏矩阵的计算效率。项目地址: https://gitcode.com/cann/ops-sparse算子概述SpMMSparse 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 的操作。调用流程为三步法GetBufferSize查询所需 workspace 大小Preprocess对稀疏矩阵进行预处理行重排 分箱加速后续计算SpMM执行 C alpha * op(A) * op(B) beta * C包含以下接口接口名功能简述aclsparseSpMMGetBufferSize查询 SpMM 所需 workspace 大小字节aclsparseSpMMPreprocess对稀疏矩阵进行预处理行重排 分箱aclsparseSpMM执行稀疏矩阵-稠密矩阵乘法 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输入aclsparseHandle_tops-sparse 库上下文句柄携带 streamHost 内存opA输入aclsparseOperation_t稀疏矩阵 A 的操作类型仅支持ACL_SPARSE_OP_NON_TRANSPOSEHost 内存opB输入aclsparseOperation_t稠密矩阵 B 的操作类型支持ACL_SPARSE_OP_NON_TRANSPOSE或ACL_SPARSE_OP_TRANSPOSEHost 内存alpha输入const void*标量 alpha 指针类型须与 computeType 匹配。内存位置由aclsparseSetPointerMode控制Host/Device 内存matA输入aclsparseConstSpMatDescr_t稀疏矩阵 A 的描述符仅支持 CSR 格式Host 内存matB输入aclsparseConstDnMatDescr_t稠密矩阵 B 的描述符Host 内存beta输入const void*标量 beta 指针类型须与 computeType 匹配。内存位置由aclsparseSetPointerMode控制Host/Device 内存matC输入/输出aclsparseDnMatDescr_t稠密矩阵 C 的描述符Host 内存computeType输入aclDataType计算精度类型支持ACL_FLOAT或ACL_INT32Host 内存alg输入aclsparseSpMMAlg_t算法类型支持ACL_SPARSE_SPMM_ALG_DEFAULT、ACL_SPARSE_SPMM_CSR_ALG1或ACL_SPARSE_SPMM_CSR_FP32_HIGH_PRECISION_ALGHost 内存size输出size_t*输出所需 workspace 大小字节Host 内存约束说明handle 不可为 nullptrmatA、matB、matC 不可为 nullptrmatA 仅支持 CSR 格式ACL_SPARSE_FORMAT_CSRmatA 的行偏移和列索引类型必须均为ACL_SPARSE_INDEX_32I且两者类型相同matA 的索引基址仅支持ACL_SPARSE_INDEX_BASE_ZEROopA 仅支持ACL_SPARSE_OP_NON_TRANSPOSEopB 支持ACL_SPARSE_OP_NON_TRANSPOSE或ACL_SPARSE_OP_TRANSPOSE不支持共轭转置数据类型组合仅支持以下三种matAACL_FLOAT, matBACL_FLOAT, matCACL_FLOAT, computeTypeACL_FLOATmatAACL_FLOAT16, matBACL_FLOAT16, matCACL_FLOAT16, computeTypeACL_FLOATmatAACL_INT8, matBACL_INT8, matCACL_INT32, computeTypeACL_INT32维度匹配A.cols B.rowsA.rows C.rowsB.cols C.colssize 不可为 nullptr支持的稀疏格式格式支持说明CSR✅稀疏矩阵 A 支持 CSR 格式COO❌不支持CSC❌不支持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)参数说明参数名输入/输出参数类型说明handle输入aclsparseHandle_tops-sparse 库上下文句柄携带 streamHost 内存opA输入aclsparseOperation_t稀疏矩阵 A 的操作类型仅支持ACL_SPARSE_OP_NON_TRANSPOSEHost 内存opB输入aclsparseOperation_t稠密矩阵 B 的操作类型支持ACL_SPARSE_OP_NON_TRANSPOSE或ACL_SPARSE_OP_TRANSPOSEHost 内存alpha输入const void*标量 alpha 指针。内存位置由aclsparseSetPointerMode控制Host/Device 内存matA输入aclsparseConstSpMatDescr_t稀疏矩阵 A 的描述符仅支持 CSR 格式Host 内存matB输入aclsparseConstDnMatDescr_t稠密矩阵 B 的描述符Host 内存beta输入const void*标量 beta 指针。内存位置由aclsparseSetPointerMode控制Host/Device 内存matC输入/输出aclsparseDnMatDescr_t稠密矩阵 C 的描述符Host 内存computeType输入aclDataType计算精度类型支持ACL_FLOAT或ACL_INT32Host 内存alg输入aclsparseSpMMAlg_t算法类型Host 内存buffer输入void*workspace 缓冲区由 GetBufferSize 返回的大小分配Device 内存约束说明同 aclsparseSpMMGetBufferSize 的约束buffer 不可为 nullptr需按 GetBufferSize 返回的大小分配支持的稀疏格式格式支持说明CSR✅稀疏矩阵 A 支持 CSR 格式COO❌不支持CSC❌不支持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)参数说明参数名输入/输出参数类型说明handle输入aclsparseHandle_tops-sparse 库上下文句柄携带 streamHost 内存opA输入aclsparseOperation_t稀疏矩阵 A 的操作类型仅支持ACL_SPARSE_OP_NON_TRANSPOSEHost 内存opB输入aclsparseOperation_t稠密矩阵 B 的操作类型支持ACL_SPARSE_OP_NON_TRANSPOSE或ACL_SPARSE_OP_TRANSPOSEHost 内存alpha输入const void*标量 alpha 指针。内存位置由aclsparseSetPointerMode控制Host/Device 内存matA输入aclsparseConstSpMatDescr_t稀疏矩阵 A 的描述符仅支持 CSR 格式Host 内存matB输入aclsparseConstDnMatDescr_t稠密矩阵 B 的描述符Host 内存beta输入const void*标量 beta 指针。内存位置由aclsparseSetPointerMode控制Host/Device 内存matC输入/输出aclsparseDnMatDescr_t稠密矩阵 C 的描述符Host 内存computeType输入aclDataType计算精度类型支持ACL_FLOAT或ACL_INT32Host 内存alg输入aclsparseSpMMAlg_t算法类型Host 内存buffer输入void*workspace 缓冲区由 GetBufferSize 返回的大小分配Device 内存约束说明同 aclsparseSpMMGetBufferSize 的约束buffer 不可为 nullptr需按 GetBufferSize 返回的大小分配调用前须先调用aclsparseSpMMPreprocess进行预处理ACL_SPARSE_SPMM_CSR_FP32_HIGH_PRECISION_ALG仅在 computeTypeACL_FLOAT 时生效使用 Kahan 补偿求和提升精度对 fp16/int8 静默忽略支持的稀疏格式格式支持说明CSR✅稀疏矩阵 A 支持 CSR 格式COO❌不支持CSC❌不支持调用示例示例代码如下仅供参考具体编译和执行过程请参考编译与运行样例。#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_ptraclsparseContext, aclsparseStatus_t (*)(aclsparseHandle_t) handlePtr(rawHandle, aclsparseDestroy); sparseRet aclsparseSetStream(static_castaclsparseHandle_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_castaclsparseHandle_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, nnz5): 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::vectorint hRowPtrA {0, 2, 3, 4, 5}; std::vectorint hColIndA {0, 2, 1, 0, 2}; std::vectorfloat 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::vectorfloat hB(static_castsize_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::vectorfloat hC(static_castsize_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_castsize_t(k) * n * sizeof(float)); void *dC AllocAndCopyDevice(hC.data(), static_castsize_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_castaclsparseHandle_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_castaclsparseHandle_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_castaclsparseHandle_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_castsize_t(m) * n * sizeof(float), dC, static_castsize_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_castlong 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创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考