描述

本 PR 完成 CANN 社区任务“08-12-aclsparseSpMM 算子开发（A2/A3）”的算子功能扩展，完善 Ascend A2、Ascend A3 平台上的 aclsparseSpMM 能力，并增加 PyTorch/ATen NPU torch.sparse.addmm 适配。

关联设计文档：

https://gitcode.com/cann/cann-ops-competitions/pull/1113

源代码信息：

• 个人仓库：https://gitcode.com/nannan-2026/ops-sparse
• 开发分支：feature/aclsparse-spmm-a2a3
• 验证提交：f385698d7c0b25890f7c3ef6acf24d5b3b933b86

改动原因

现有 arch22 aclsparseSpMM 实现的数据类型、矩阵布局、转置方式、稀疏边界处理以及 PyTorch 调用链路尚不完整，无法完整覆盖任务书要求的 A2/A3 功能和测试场景。

本 PR 在现有 ops-sparse 目录和接口基础上进行扩展，避免重复创建同名算子，主要解决以下问题：

1. 数据类型支持范围不完整；
2. CSR 索引基址、矩阵布局和算法参数处理不完整；
3. opA、opB 转置及共轭转置场景不完整；
4. nnz=0、空行等 CSR 边界场景缺少完整处理；
5. 缺少 PyTorch/ATen NPU torch.sparse.addmm 到 aclsparseSpMM 的调用适配；
6. 测试程序一次执行全部用例耗时较长，缺少单用例选择和专项 smoke 测试能力；
7. ATK worker 子进程无法自动完成自定义 ATen NPU 扩展注册。

主要改动

1. 扩展 aclsparseSpMM 接口能力

完善 include/cann_ops_sparse.h 及 arch22 SpMM Host、Kernel 侧实现，支持：

• FP16；
• BF16；
• FP32；
• Complex64；
• CSR int32 索引；
• CSR index base 0；
• CSR index base 1；
• Row Major 稠密矩阵；
• Column Major 稠密矩阵；
• 非默认 leading dimension；
• SpMM 算法参数选择；
• opA 非转置、转置和共轭转置；
• opB 非转置、转置和共轭转置；
• alpha、beta 缩放；
• beta=0 场景；
• 稠密 input 广播；
• 非连续稠密 Tensor 输入；
• nnz=0；
• 中间空行；
• CSR 首尾空行及其他稀疏边界场景。

2. 完善 Host 侧校验与 Tiling

Host 侧增加或完善以下逻辑：

• 输入描述符和数据类型检查；
• CSR row pointer、column index、value 信息处理；
• 矩阵维度及转置后维度推导；
• Row Major、Column Major 地址布局处理；
• leading dimension 合法性检查；
• CSR index base 处理；
• 算法参数传递；
• Workspace 大小计算；
• 稀疏边界条件处理；
• Host 到 Kernel 的 Tiling 参数传递。

3. 扩展 Ascend C Kernel

Kernel 侧增加或完善以下能力：

• 不同数据类型的读取、转换和计算；
• FP16/BF16 中间计算处理；
• Complex64 实部、虚部乘加；
• 共轭转置处理；
• Row Major、Column Major 地址计算；
• CSR base 0/base 1 索引归一化；
• alpha * op(A) * op(B) + beta * C 计算；
• beta=0 时避免读取无效或非有限输入；
• nnz=0、空行等边界处理；
• 输出布局和 leading dimension 写回。

4. 增加 PyTorch/ATen NPU Adapter

新增 python_adapter，通过 PyTorch PrivateUse1/SparseCsrPrivateUse1 调度机制，将：

torch.sparse.addmm(input, mat1, mat2, beta, alpha)

映射到：

aclsparseSpMM

Adapter 主要包含：

• CSR Tensor 参数提取；
• 稠密 Tensor 参数提取；
• dtype 和 layout 转换；
• 输入广播；
• 非连续 Tensor 连续化处理；
• Workspace 申请；
• ACL Stream 获取与执行；
• ACL 错误码检查；
• 自定义 ATen NPU Kernel 注册；
• FP16、BF16、FP32、Complex64 支持。

同时增加：

• Adapter 构建脚本；
• Adapter 独立测试脚本；
• sitecustomize.py 自动加载机制；
• OPS_SPARSE_TORCH_LIBRARY 环境变量支持；
• ATK worker 子进程自动注册能力；
• Adapter 使用说明文档。

5. 完善 arch22 SpMM 测试

扩展 test/spmm/arch22/spmm_test.cpp，增加：

• SPMM_CASE_INDEX 单用例选择；
• 非法用例索引校验；
• FP16/BF16 dtype smoke；
• layout/base/algorithm smoke；
• Complex64 smoke；
• opA、opB 转置及共轭转置测试；
• nnz=0 边界测试；
• 中间空行测试；
• alpha、beta 组合测试；
• 精度差异统计和明确的退出状态。

测试程序支持按场景快速回归，避免每次调试都执行全部用例。

关联的Issue

无独立 Issue。

关联社区任务：

08-12-aclsparseSpMM 算子开发（A2/A3）

关联设计文档 PR：

https://gitcode.com/cann/cann-ops-competitions/pull/1113

测试

本 PR 已在真实 Ascend A2、Ascend A3 环境中进行构建、功能、精度和调用链路验证。

A2 环境验证

• CANN：9.0.0；
• 代码分支：feature/aclsparse-spmm-a2a3；
• 代码提交：f385698d7c0b25890f7c3ef6acf24d5b3b933b86；
• 工程构建：通过；
• libops_sparse.so 安装：通过；
• ATen NPU Adapter 构建：通过；
• PyTorch torch.sparse.addmm 调用：通过。

专项测试结果：

• FP16 dtype smoke：通过；
• BF16 dtype smoke：通过；
• layout/base/algorithm smoke：通过；
• Complex64 smoke：通过；
• opA 转置及共轭转置：通过；
• opB 转置及共轭转置：通过；
• nnz=0：通过；
• 中间空行：通过；
• 代表性 SpMM 用例：通过；
• ATen NPU FP16：通过；
• ATen NPU BF16：通过；
• ATen NPU FP32：通过；
• ATen NPU Complex64：通过。

A2 ATK 完整精度测试结果：

Total Task: 200, success 200, failed 0
Summary info: acc_pass_result:Pass
atk task success

A2 精度通过率：100%。

A3 环境验证

• CANN：9.0.0；
• PyTorch：2.7.1+cpu；
• torch_npu：2.7.1.post4；
• NPU：Ascend A3 环境；
• 代码分支：feature/aclsparse-spmm-a2a3；
• 代码提交：f385698d7c0b25890f7c3ef6acf24d5b3b933b86；
• 工程构建：通过；
• ATen NPU Adapter 强制清理重建：通过；
• PyTorch torch.sparse.addmm 调用：通过。

专项测试结果：

• FP16 dtype smoke：通过；
• BF16 dtype smoke：通过；
• layout/base/algorithm smoke：通过；
• Complex64 smoke：通过；
• opA、opB 转置及共轭转置：通过；
• nnz=0：通过；
• 中间空行：通过；
• ATen NPU FP16：通过；
• ATen NPU BF16：通过；
• ATen NPU FP32：通过；
• ATen NPU Complex64：通过。

A3 ATK 完整精度测试结果：

Total Task: 200, success 200, failed 0
Summary info: acc_pass_result:Pass
atk task success

A3 精度通过率：100%。

A3 性能采集链路预检结果：

Event performance: completed 2/2 succeeded
EVENT_SMOKE_EXIT=0
Profiler Level1: completed 2/2 succeeded
PROFILER_LEVEL1_EXIT=0

Profiler 已成功生成：

• op_statistic.csv；
• kernel_details.csv；
• NPU Profiler 汇总 CSV；
• NPU Profiler 汇总 JSON。

上述结果用于确认 torch.sparse.addmm 调用实际进入 NPU Kernel。当前 PR 不声明尚未形成正式完整报告的性能结论。

构建与代码检查

• git diff --check：通过；
• CMake 构建：通过；
• Ascend C Kernel 编译：通过；
• C++ 测试程序编译：通过；
• Python Adapter 编译：通过；
• GitCode Hooks：通过；
• 工作区状态检查：通过；
• 远端分支同步：通过；
• 无账号、令牌、密码等敏感信息。

文档更新

本 PR 新增和更新了以下文档内容：

• 新增 python_adapter/README.md；
• 说明 ATen NPU Adapter 的构建方法；
• 说明 OPS_SPARSE_TORCH_LIBRARY 环境变量；
• 说明 sitecustomize.py 自动加载方法；
• 说明 torch.sparse.addmm 的调用和验证方法；
• 补充关联设计文档 PR。

完整算子设计文档位于：

https://gitcode.com/cann/cann-ops-competitions/pull/1113

类型标签

• Bug修复
• 新特性
• 性能优化
• 文档更新
• 其他，请描述：新增测试用例、PyTorch/ATen NPU Adapter 及 A2/A3 验证支持