[论文解读] SparseTIR: Composable Abstractions for Sparse Compilation in Deep Learning

Zihao Ye, Ruihang Lai+3 more|arXiv (Cornell University)|Jul 11, 2022

Advanced Neural Network Applications被引用 5

一句话总结

SparseTIR 引入了一种用于稀疏深度学习编译的可组合中间表示（IR），通过可组合的格式和转换实现高效优化。它在稀疏算子上相比厂商库实现了 1.05–7.45× 的加速，在端到端图神经网络（GNN）推理中最高实现 40.18× 的加速，通过动态组合最优稀疏格式和硬件感知优化实现。

ABSTRACT

Sparse tensors are rapidly becoming critical components of modern deep learning workloads. However, developing high-performance sparse operators can be difficult and tedious, and existing vendor libraries cannot satisfy the escalating demands from new operators. Sparse tensor compilers simplify the development of operators, but efficient sparse compilation for deep learning remains challenging because a single sparse format cannot maximize hardware efficiency, and single-shot compilers cannot keep up with latest hardware and system advances. In this paper, we observe that the key to addressing both these challenges is to leverage composable formats and composable transformations. We propose SparseTIR, a sparse tensor compilation abstraction that offers composable formats and composable transformations for deep learning workloads. SparseTIR constructs a search space over these composable components for performance tuning. With these improvements, SparseTIR obtains consistent performance speedups vs vendor libraries on GPUs for single operators: 1.20-2.34x for GNN operators, 1.05-2.98x for sparse attention operators, and 0.56-7.45x for sparse convolution operators. SparseTIR also accelerates end-to-end GNNs by 1.08-1.52x for GraphSAGE training, and 4.20-40.18x for RGCN inference.

研究动机与目标

• 解决深度学习中高性能稀疏算子开发的挑战，现有厂商库无法支持新兴且多样的工作负载。
• 克服单一格式稀疏编译器无法适应异构硬件或不断演化的深度学习工作负载的局限。
• 通过解耦格式与调度决策，实现稀疏算子的高效、自动化性能调优。
• 通过可组合的编译抽象，支持复杂深度学习工作负载（如 GNN 和稀疏 Transformer）的端到端加速。
• 提供可扩展、可扩展的框架，在复用现有密集张量编译器优化的同时，支持新的稀疏专用转换。

提出的方法

• 提出 SparseTIR，一种可组合的中间表示（IR），将稀疏格式规范与计算描述分离。
• 通过允许稀疏张量的不同部分基于局部稀疏模式存储在不同最优格式中，实现格式的可组合性。
• 通过将编译过程分解为一系列可重用、可组合的程序转换，实现转换的可组合性。
• 复用密集张量编译器的现有循环级抽象（如张量化、分块），并扩展其以支持稀疏数据处理。
• 构建一个覆盖可组合格式与转换的搜索空间，通过可配置的搜索系统指导自动化性能调优。
• 集成来自现有编译器的硬件特定优化（如张量核心、GPU 映射），以加速现代加速器上的稀疏内核。

实验结果

研究问题

• RQ1 可组合格式是否能提升在多样化稀疏深度学习工作负载中的性能可移植性和硬件利用率？
• RQ2 可组合转换是否能实现密集张量编译器优化在稀疏编译中的高效复用？
• RQ3 基于可组合组件的搜索调优系统是否能在新兴稀疏算子上超越厂商优化库的性能？
• RQ4 SparseTIR 相较基线库，能在多大程度上加速端到端 GNN 训练与推理？
• RQ5 格式可组合性如何促进稀疏内核对硬件特性（如张量核心和内存层次结构）的更好利用？

主要发现

• 在 GNN SpMM 算子上，SparseTIR 相比厂商库实现了 1.20–2.34× 的加速，尤其在高度不规则的稀疏模式下收益最大。
• 在稀疏注意力算子上，SparseTIR 实现了 1.05–2.98× 的加速，表明其在中等稀疏、结构化稀疏场景下的有效性。
• 在稀疏卷积算子上，SparseTIR 最高实现 7.45× 的加速，在基于点云的工作负载中显著优于厂商库。
• 端到端 GraphSAGE 训练加速了 1.08–1.52×，表明在完整模型训练流水线中具有持续的性能增益。
• RGCN 推理获得显著提升，加速比在 4.20× 到 40.18× 之间，凸显了在复杂、异构稀疏性场景中可组合优化的优势。
• 该成果在不同 GPU（RTX 3070、V100）上均表现出可复现性，当启用 L2 缓存清除时，结果与论文报告数据高度一致。