Algorithm to Compilation Co-design: An Integrated View of Neural Network Sparsity
This work addresses the gap between algorithm-level sparsity and hardware execution for practitioners, though it is incremental as it builds on existing pruning and compiler techniques.
The paper tackles the problem of translating neural network sparsity into actual runtime performance gains by co-designing pruning algorithms and compiler support, finding that structured sparsity with block sparse representation (BSR) in TVM yields up to 4x speedup over vanilla PyTorch and 2.2x over standard TVM, with optimal block shapes being linear rather than square for BERT attention weights on CPU.
Reducing computation cost, inference latency, and memory footprint of neural networks are frequently cited as research motivations for pruning and sparsity. However, operationalizing those benefits and understanding the end-to-end effect of algorithm design and regularization on the runtime execution is not often examined in depth. Here we apply structured and unstructured pruning to attention weights of transformer blocks of the BERT language model, while also expanding block sparse representation (BSR) operations in the TVM compiler. Integration of BSR operations enables the TVM runtime execution to leverage structured pattern sparsity induced by model regularization. This integrated view of pruning algorithms enables us to study relationships between modeling decisions and their direct impact on sparsity-enhanced execution. Our main findings are: 1) we validate that performance benefits of structured sparsity block regularization must be enabled by the BSR augmentations to TVM, with 4x speedup relative to vanilla PyTorch and 2.2x speedup relative to standard TVM compilation (without expanded BSR support). 2) for BERT attention weights, the end-to-end optimal block sparsity shape in this CPU inference context is not a square block (as in \cite{gray2017gpu}) but rather a linear 32x1 block 3) the relationship between performance and block size / shape is is suggestive of how model regularization parameters interact with task scheduler optimizations resulting in the observed end-to-end performance.