资讯动态

SpMM 从哪儿开始:把 CSR 的行边界和输出语义写清楚

发布时间:2026/10/25 14:10:12 来源:尧图企业网站定制
稀疏矩阵乘法不等于把 GEMM 里的零跳过去。索引、分支和不规则读取也要花时间块里只有少数非零元素时BCSR 还可能把许多零重新带回来。先决定数据怎么表示、输出是覆盖还是累加比先报一个稀疏率更有用。稀疏 A 为 M×K稠密 B 为 K×N。CSR 用row_ptr划出每行在values和col_idx中的区间。本例约定同一位置出现两次就相加结果从零开始所以实现的是CA×B。如果直接往已有 C 上那是另一个接口。CSR 将每行边界、列索引和值分开保存。def spmm(ptr,cols,vals,b,k,n): if not ptr or ptr[0]!0 or ptr[-1]!len(vals): raise ValueError(invalid row pointer) if len(cols)!len(vals) or len(b)!k*n: raise ValueError(storage mismatch) if any(xy for x,y in zip(ptr,ptr[1:])): raise ValueError(non-monotone row pointer) if any(not 0ck for c in cols): raise ValueError(column out of bounds) out[0.] * ((len(ptr)-1)*n) for i in range(len(ptr)-1): for p in range(ptr[i],ptr[i1]): for j in range(n): out[i*nj]vals[p]*b[cols[p]*nj] return out # A[[2,0,3],[0,0,0]], B[[1,2],[3,4],[5,6]] assert spmm([0,2,2],[0,2],[2.,3.],[1.,2.,3.,4.,5.,6.],3,2)[17.,22.,0.,0.] # Duplicate coordinates have additive semantics. assert spmm([0,2],[0,0],[1.,2.],[4.],1,1)[12.] print(CSR: empty row and duplicate coordinate passed)稀疏率之外还要看分布一行一个非零和一行上千个非零即使整体密度相同任务分配也完全不同。按行平分线程可能让某个线程拖住全局。N 很小时索引成本显眼N 大一些同一个非零 A 值可以复用于 B 的多列。转成 BCSR 之前可以统计每个候选块的占用率估算为块内零支付的载荷。2026 年 9 月的 PyTorch 2.14 发布说明介绍了 NVGEMM 等后端进展。它提醒我们新的稠密或块缩放路径需要放回实际算子里比较。非结构化稀疏不能自动套用结构化硬件稀疏的加速承诺也不能仅靠减少乘法次数推导耗时。接到系统里的边界Rust 中可以把 CSR 构造时的校验集中到一个类型里让后续计算只接收有效表示。但动态更新仍可能破坏已建立的条件接口要么禁止更新要么重新验证。usize的维度乘法还需checked_mul否则长度检查本身也可能溢出。用稀疏结构本身解释一次结果示例第一行只保存第零列和第二列所以输出第一行应是 B 第零行的两倍加第二行的三倍。第二行没有任何非零项因此必须完整写出一行零。这个空行很重要若实现跳过整行却沿用旧输出第一次调用可能正常复用缓冲区后就会留下上一请求的数据。重复坐标也需要固定语义。这里选择求和允许装配阶段把多个贡献先写成重复项另一些接口可能要求预先合并。若要缓存排序后的 CSR就应把排序和合并规则视为格式的一部分。仅校验列号范围无法判断输入是不是满足某个 kernel 的额外排序约束。行指针可以有相邻相等的值这恰好表示空行不能误判成非法。它还必须单调、不超出载荷末尾并以非零总数结束。本例在元素个数一致的前提下用这些约束形成索引边界生产解析器还应验证字段类型与整数范围避免负数或过大值进入底层索引。讨论 BCSR 时可以先只做一次离线统计候选块中有多少有效非零块索引节省多少补零又引入多少乘法。这份账能够筛掉明显不合适的块形状但不能预言设备耗时。若决定转换保留转换前后矩阵的逐元素参考避免索引重排错误被性能测量掩盖。稀疏格式的排序还可能影响浮点相加次序。重复项合并前后得到很小差异时应先检查数值范围和累加精度若某一整行结果不同则更可能是索引或清零合同。把两类问题分开可以避免用误差阈值掩盖结构错误也能避免把正常舍入差异当成格式损坏。这里实测的是小规模 CSR 的代数结果没有测线程、向量化、缓存和模型权重。评估优化时应另外记录非零分布、索引位宽、转换成本及稠密参考误差格式转换若只用一次它可能比省下的计算更贵。

读完文章,也想定制专属网站?

尧图设计师 24 小时内与您沟通定制方案

免费获取报价 →
↑