一句话:提出稀疏寄存器分块,通过 unroll-and-sparse-jam 与数据压缩,为稀疏神经网络推理的 SpMM 内核提升寄存器复用并加速。
非结构化稀疏神经网络可压缩模型并减少浮点运算,但其推理常受稀疏矩阵乘 SpMM(C=A×B)内核主导。剪枝模型的稀疏模式与稀疏度使既有稀疏库优化失效,且因访存地址难以静态确定,寄存器复用困难。本文提出稀疏寄存器分块(Sparse Register Tiling),由面向 ML 矩阵稀疏模式的 unroll-and-sparse-jam 变换与数据压缩组成,利用稀疏信息 jamming 代码以提升寄存器复用。在 2396 个 transformer 与卷积模型权重矩阵、60–95% 稀疏度上评估,多核 CPU 上较 MKL SpMM 平均加速 1.72×、较稠密矩阵乘加速 2.65×;ARM 边缘设备上 MobileNetV1(70% 稀疏)端到端加速 2.12×。
Unstructured sparse neural networks are an important class of machine learning (ML) models, as they compact model size and reduce floating point operations. The execution time of these models is frequently dominated by the sparse matrix multiplication (SpMM) kernel, C = A × B , where A is a sparse matrix, and B and C are dense matrices. The unstructured sparsity pattern of matrices in pruned machine learning models along with their sparsity ratio has rendered useless the large class of libraries and systems that optimize sparse matrix multiplications. Reusing registers is particularly difficult because accesses to memory locations should be known statically. This paper proposes Sparse Register Tiling, a new technique composed of an unroll-and-sparse-jam transformation followed by data compression that is specifically tailored to sparsity patterns in ML matrices. Unroll-and-sparse-jam uses sparsity information to jam the code while improving register reuse. Sparse register tiling is evaluated across 2396 weight matrices from transformer and convolutional models with a sparsity range of 60-95% and provides an average speedup of 1.72× and 2.65× over MKL SpMM and dense matrix multiplication, respectively, on a multicore CPU processor. It also provides an end-to-end speedup of 2.12× for MobileNetV1 with 70% sparsity on an ARM processor commonly used in edge devices.
@article{paperbot1957,
title = {Register Tiling for Unstructured Sparsity in Neural Network Inference},
author = {Lucas Wilkinson and Kazem Cheshmi and Maryam Mehri Dehnavi},
journal = {Proceedings of the ACM on Programming Languages},
volume = {7},
number = {PLDI},
year = {2023},
doi = {10.1145/3591302}
}