CANN ops-sparse 非零元素统计
非零元素统计算子【免费下载链接】ops-sparse本项目是CANN提供的高性能稀疏矩阵计算的算子库专注于优化稀疏矩阵的计算效率。项目地址: https://gitcode.com/cann/ops-sparse算子概述nnz 算子用于统计稠密矩阵中每行或每列的非零元素个数以及总非零元素个数。该函数是 dense-to-sparse 格式转换流程的第一步典型用法为先调用 nnz 获取每行/列非零元素个数和总非零元素个数再根据统计结果构建 CSR 或 CSC 格式的偏移数组通过前缀和最后执行 dense-to-sparse 转换。核心运算为dirA ROW: nnzPerRow[i] sum_{j0}^{n-1} [A[i][j] ! 0], i in [0, m) dirA COLUMN: nnzPerCol[j] sum_{i0}^{m-1} [A[i][j] ! 0], j in [0, n) nnzTotal sum(nnzPerRow) sum(nnzPerCol)其中[A[i][j] ! 0]为指示函数当A[i][j] ! 0时取 1否则取 0。包含以下接口接口名功能简述aclsparseSnnz统计 FP32 稠密矩阵每行或每列的非零元素个数及总非零元素个数算子执行接口aclsparseSnnz产品支持情况Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品不支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品不支持函数原型aclsparseStatus_t aclsparseSnnz(aclsparseHandle_t handle, aclsparseDirection_t dirA, int m, int n, const aclsparseMatDescr_t descrA, const float *A, int lda, int *nnzPerRowColumn, int *nnzTotalDevHostPtr)参数说明参数名输入/输出参数类型说明handle输入aclsparseHandle_tops-sparse 库上下文句柄携带 streamHost 内存dirA输入aclsparseDirection_t统计方向ACL_SPARSE_DIRECTION_ROW按行或ACL_SPARSE_DIRECTION_COLUMN按列Host 内存m输入int矩阵 A 的行数Host 内存n输入int矩阵 A 的列数Host 内存descrA输入const aclsparseMatDescr_t矩阵描述符仅支持ACL_SPARSE_MATRIX_TYPE_GENERAL类型Host 内存A输入const float*FP32稠密矩阵数据列主序存储维度为 lda × n元素A[i][j] A[j * lda i]Device 内存lda输入int矩阵 A 的 leading dimensionHost 内存nnzPerRowColumn输出int*每行dirAROW 时大小为 m或每列dirACOLUMN 时大小为 n的非零元素个数数组Device 内存nnzTotalDevHostPtr输出int*总非零元素个数内存位置由aclsparseSetPointerMode控制HOST 模式为 Host 内存DEVICE 模式为 Device 内存约束说明m 0n 0lda max(1, m)descrA 不可为 nullptr且 descrA.type 仅支持ACL_SPARSE_MATRIX_TYPE_GENERAL当 m 0 且 n 0 时A 不可为 nullptr当 totalUnits 0 时totalUnits m 或 n取决于 dirAnnzPerRowColumn 不可为 nullptrnnzTotalDevHostPtr 不可为 nullptr当 m 0 或 n 0 时函数直接返回成功nnzTotal 置 0nnzPerRowColumn 不写入调用示例示例代码如下仅供参考具体编译和执行过程请参考编译与运行样例。#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; }; int aclsparseSnnzTest(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); // 使用 DEVICE 模式nnzTotalDevHostPtr 指向 Device 内存 sparseRet aclsparseSetPointerMode(static_castaclsparseHandle_t(handlePtr.get()), ACL_SPARSE_POINTER_MODE_DEVICE); CHECK_RET(sparseRet ACL_SPARSE_STATUS_SUCCESS, LOG_PRINT(aclsparseSetPointerMode failed. ERROR: %d\n, sparseRet); return sparseRet); // 2. 准备 Host 数据列主序 3x3 矩阵 // 矩阵内容 // [1.0 0.0 3.0] // [0.0 0.0 4.0] // [2.0 0.0 0.0] int m 3, n 3, lda 3; std::vectorfloat hA(static_castsize_t(lda) * n, 0.0f); hA[0 * lda 0] 1.0f; // A[0][0] hA[0 * lda 2] 2.0f; // A[2][0] hA[2 * lda 0] 3.0f; // A[0][2] hA[2 * lda 1] 4.0f; // A[1][2] // 3. 申请 Device 内存并拷贝数据 void* rawMemA nullptr; auto aclRet aclrtMalloc(rawMemA, static_castsize_t(lda) * n * sizeof(float), ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(aclRet ACL_SUCCESS, LOG_PRINT(aclrtMalloc for A failed. ERROR: %d\n, aclRet); return aclRet); std::unique_ptrvoid, aclError (*)(void*) dA(rawMemA, aclrtFree); aclRet aclrtMemcpy(dA.get(), static_castsize_t(lda) * n * sizeof(float), hA.data(), static_castsize_t(lda) * n * sizeof(float), ACL_MEMCPY_HOST_TO_DEVICE); CHECK_RET(aclRet ACL_SUCCESS, LOG_PRINT(aclrtMemcpy for A failed. ERROR: %d\n, aclRet); return aclRet); void* rawMemNnzPerRow nullptr; aclRet aclrtMalloc(rawMemNnzPerRow, m * sizeof(int), ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(aclRet ACL_SUCCESS, LOG_PRINT(aclrtMalloc for nnzPerRow failed. ERROR: %d\n, aclRet); return aclRet); std::unique_ptrvoid, aclError (*)(void*) dNnzPerRow(rawMemNnzPerRow, aclrtFree); void* rawMemNnzTotal nullptr; aclRet aclrtMalloc(rawMemNnzTotal, sizeof(int), ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(aclRet ACL_SUCCESS, LOG_PRINT(aclrtMalloc for nnzTotal failed. ERROR: %d\n, aclRet); return aclRet); std::unique_ptrvoid, aclError (*)(void*) dNnzTotal(rawMemNnzTotal, aclrtFree); // 4. 创建 MatDescrLegacy API aclsparseMatDescr_t matDescr nullptr; sparseRet aclsparseCreateMatDescr(matDescr); CHECK_RET(sparseRet ACL_SPARSE_STATUS_SUCCESS, LOG_PRINT(aclsparseCreateMatDescr failed. ERROR: %d\n, sparseRet); return sparseRet); aclsparseSetMatType(matDescr, ACL_SPARSE_MATRIX_TYPE_GENERAL); aclsparseSetMatIndexBase(matDescr, ACL_SPARSE_INDEX_BASE_ZERO); // 5. 调用 aclsparseSnnz按行统计 aclsparseDirection_t dirA ACL_SPARSE_DIRECTION_ROW; sparseRet aclsparseSnnz( static_castaclsparseHandle_t(handlePtr.get()), dirA, m, n, matDescr, static_castconst float*(dA.get()), lda, static_castint*(dNnzPerRow.get()), static_castint*(dNnzTotal.get())); CHECK_RET(sparseRet ACL_SPARSE_STATUS_SUCCESS, LOG_PRINT(aclsparseSnnz failed. ERROR: %d\n, sparseRet); aclsparseDestroyMatDescr(matDescr); return sparseRet); // 6. 同步等待任务执行结束 aclRet aclrtSynchronizeStream(stream); CHECK_RET(aclRet ACL_SUCCESS, LOG_PRINT(aclrtSynchronizeStream failed. ERROR: %d\n, aclRet); return aclRet); // 7. 将结果从 Device 拷贝回 Host 并打印 std::vectorint hNnzPerRow(m, 0); aclRet aclrtMemcpy(hNnzPerRow.data(), m * sizeof(int), dNnzPerRow.get(), m * sizeof(int), ACL_MEMCPY_DEVICE_TO_HOST); CHECK_RET(aclRet ACL_SUCCESS, LOG_PRINT(copy nnzPerRow from device to host failed. ERROR: %d\n, aclRet); return aclRet); int hNnzTotal 0; aclRet aclrtMemcpy(hNnzTotal, sizeof(int), dNnzTotal.get(), sizeof(int), ACL_MEMCPY_DEVICE_TO_HOST); CHECK_RET(aclRet ACL_SUCCESS, LOG_PRINT(copy nnzTotal from device to host failed. ERROR: %d\n, aclRet); return aclRet); for (int i 0; i m; i) { LOG_PRINT(nnzPerRow[%d] is: %d\n, i, hNnzPerRow[i]); } LOG_PRINT(nnzTotal is: %d\n, hNnzTotal); // 8. 清理 MatDescr aclsparseDestroyMatDescr(matDescr); return ACL_SPARSE_STATUS_SUCCESS; } int main() { AclContext ctx(0); auto ret ctx.Init(); CHECK_RET(ret ACL_SUCCESS, return ret); ret aclsparseSnnzTest(ctx); CHECK_RET(ret ACL_SPARSE_STATUS_SUCCESS, LOG_PRINT(aclsparseSnnzTest failed. ERROR: %d\n, ret); return ret); return 0; }预期输出如下nnzPerRow[0] is: 2 nnzPerRow[1] is: 1 nnzPerRow[2] is: 1 nnzTotal is: 4【免费下载链接】ops-sparse本项目是CANN提供的高性能稀疏矩阵计算的算子库专注于优化稀疏矩阵的计算效率。项目地址: https://gitcode.com/cann/ops-sparse创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考