本文介绍一种兼顾性能与兼容性的稀疏矩阵乘法加速方案：在 Numba nopython 模式下，通过 objmode 安全调用 SciPy 高度优化的稀疏乘法内核，并直接返回 COO 或 CSR 格式结果，实测仅比原生 SciPy 慢 2%–15%。
本文介绍一种兼顾性能与兼容性的稀疏矩阵乘法加速方案：在 numba `nopython` 模式下，通过 `objmode` 安全调用 scipy 高度优化的稀疏乘法内核，并直接返回 coo 或 csr 格式结果，实测仅比原生 scipy 慢 2%–15%。
在高性能科学计算中，稀疏矩阵乘法（SpMM）是常见瓶颈。尽管 Numba 不原生支持 scipy.sparse 类型，且手动实现 COO 格式乘法（如双循环 + 哈希累加）看似可控，但实际性能往往远逊于 SciPy —— 这源于 SciPy 底层调用高度优化的 CSR/CSC 内核（基于 Intel MKL、OpenBLAS 或自研稀疏算法），并经过数十年工程打磨。
直接手写 Numba 版本（如问题中 _mul 函数）面临多重挑战：
- 内存局部性差：COO 三元组无序存储，导致频繁随机访存；
- 重复索引合并开销大：需动态哈希或排序去重，而 np.zeros((an, bm)) 的稠密中间矩阵更违背稀疏初衷；
- 并行化失效：@njit(parallel=True) 在细粒度稀疏访存场景下易引发线程竞争与缓存抖动，反而降速；
- 内存预分配困难：输出非零元数量未知，extend_arr 动态扩容带来显著额外开销。
因此，最优策略不是重造轮子，而是“桥接”——在 Numba 生态中安全复用 SciPy 的工业级实现。关键在于：如何在保持调用方 @njit 兼容性的同时，嵌入 Python 层稀疏运算？
答案是 numba.objmode：它允许在 nopython 函数中划定一段“Python 模式”代码块，执行任意 Python 对象操作（如构建 coo_matrix、调用 @ 运算符、转换格式），同时严格声明该块的输入/输出类型，使 Numba 能静态推断整个函数签名。
以下为两个生产就绪的实现：
✅ 推荐方案 1：COO 格式输出（兼顾通用性）
import numba as nb
from scipy.sparse import coo_matrix
@nb.njit()
def mul_coo(ar, ac, av, br, bc, bv, n):
"""
稀疏矩阵乘法：A @ B，输入为 COO 三元组，输出为 (row, col, data) 三元组。
A = coo_matrix((av, (ar, ac)), shape=(n, n))
B = coo_matrix((bv, (br, bc)), shape=(n, n))
"""
with nb.objmode(row='i4[:]', col='i4[:]', data='f8[:]'):
a_sci = coo_matrix((av, (ar, ac)), shape=(n, n))
b_sci = coo_matrix((bv, (br, bc)), shape=(n, n))
res_coo = (a_sci @ b_sci).tocoo() # 强制转为 COO
row = res_coo.row.copy() # .row/.col 是 view，需 copy 保证所有权
col = res_coo.col.copy()
data = res_coo.data.copy()
return row, col, data✅ 推荐方案 2：CSR 格式输出（极致性能）
@nb.njit()
def mul_csr(ar, ac, av, br, bc, bv, n):
"""
输出 CSR 格式三元组：(data, indices, indptr)，避免 tocoo() 开销。
更适合后续 CSR 专用计算（如 SpMV）。
"""
with nb.objmode(data='f8[:]', indices='i4[:]', indptr='i4[:]'):
a_sci = coo_matrix((av, (ar, ac)), shape=(n, n))
b_sci = coo_matrix((bv, (br, bc)), shape=(n, n))
res_csr = a_sci @ b_sci # 默认返回 CSR
data = res_csr.data.copy()
indices = res_csr.indices.copy()
indptr = res_csr.indptr.copy()
return data, indices, indptr⚠️ 关键注意事项
- copy() 不可省略：res_coo.row 等是 NumPy view，若不显式 copy()，Numba 可能因内存所有权问题报错或产生未定义行为；
- objmode 块内禁止 Numba 类型操作：所有 SciPy 构建、乘法、格式转换必须在 with nb.objmode(...): 内完成；
- 类型声明须精确：'i4[:]' 表示 int32 一维数组，'f8[:]' 表示 float64 一维数组，与 NumPy dtype 严格对应；
- 避免过度使用：objmode 会中断 JIT 流水线，仅用于不可替代的 Python 生态调用；若全程需纯 nopython，应转向 CSR/CSC 手写内核（但开发成本与调试难度剧增）；
- 形状一致性：示例假设方阵 (n, n)，实际使用时请按需传入 shape=(m, k) 和 (k, p) 并校验维度兼容性。
? 性能对比（n=50000, m=1000）
| 方法 | 耗时（均值 ± std） | 相对 SciPy |
|---|---|---|
| 原生 SciPy (coo @ coo) | 27.8 ms ± 471 µs | 1.0×（基准） |
| 手写 Numba _mul | 184 ms ± 594 µs | ≈6.6× 慢 |
| mul_coo（COO 输出） | 32.0 ms ± 228 µs | ≈1.15× 慢 |
| mul_csr（CSR 输出） | 28.3 ms ± 685 µs | ≈1.02× 慢 |
可见，objmode 方案成功将性能损失控制在极小范围内，同时保留了 Numba 函数链的 nopython 兼容性——上游数据预处理、下游稀疏向量运算等均可无缝使用 @njit 加速。
总结：当面对稀疏计算这类“已有成熟工业实现”的任务时，明智的加速策略是“站在巨人的肩膀上”。numba.objmode 提供了安全、类型明确、低开销的桥梁，让 Numba 用户得以在不牺牲生态优势的前提下，获得接近底层库的性能。这正是现代高性能 Python 工程化的典型范式。
