一句话:提出 FEASTA,从数据流抽象、ISA 与架构三层协同设计,高效灵活加速机器学习中的通用稀疏张量代数。
稀疏张量代数(SpTA)在机器学习中日益重要,但其非结构化稀疏性使 GPU、CPU 等通用处理器硬件资源利用率低,而专用稀疏核加速器又难以灵活支持不同数据流与稀疏度的任务。本文提出 FEASTA,从三个层次进行设计:在数据流抽象层,将 Einstein Summation 应用于稀疏 fiber tree,把通用 SpTA 的统一执行流建模为 fiber tree 的连接与合并;在指令集架构层,基于该执行流提出通用 SpTA ISA,包含面向稠密和稀疏数据的不同指令;在架构层,设计由可配置、高性能功能单元组成的指令驱动架构。实验表明,FEASTA 相比 GPU 在多种负载上取得 5.40× 几何平均能效提升,在稀疏矩阵乘法上比最先进稀疏矩阵加速器和 CPU 扩展分别快 1.47× 与 3.19×,并在多种内核上实现 1.69–12.70× 的能效改进。
Recently, sparse tensor algebra (SpTA) plays an increasingly important role in machine learning. However, due to the unstructured sparsity of SpTA, the general-purpose processors (e.g., GPU and CPU) are inefficient because of the underutilized hardware resources. Sparse kernel accelerators are optimized for specific tasks. However, their dedicated processing units and data paths cannot effectively support other SpTA tasks with different dataflow and various sparsity, resulting in performance degradation. This paper proposes FEASTA, a Flexible and Efficient Accelerator for Sparse Tensor Algebra. To process general SpTA tasks with various sparsity efficiently, we design FEASTA meticulously from three levels. At the dataflow abstraction level, we apply the Einstein Summation on the sparse fiber tree data structure to model the unified execution flow of general SpTA as joining and merging the fiber tree. At the instruction set architecture (ISA) level, a general SpTA ISA is proposed based on the execution flow. It includes different types of instructions for dense and sparse data, achieving flexibility and efficiency at the instruction level. At the architecture level, an instruction-driven architecture consisting of configurable and high-performance function units is designed, supporting the flexible and efficient ISA. Evaluations show that FEASTA has 5.40× geomean energy efficiency improvements compared to GPU among various workloads. FEASTA delivers 1.47× and 3.19× higher performance on sparse matrix multiplication kernels compared to state-of-the-art sparse matrix accelerator and CPU extension. Across diverse kernels, FEASTA achieves 1.69-12.70× energy efficiency over existing architectures.
@inproceedings{ZhongZDWYZSMZHZ24,
title = {FEASTA: A Flexible and Efficient Accelerator for Sparse Tensor Algebra in Machine Learning},
author = {Kai Zhong and Zhenhua Zhu and Guohao Dai and Hongyi Wang and Xinhao Yang and Haoyu Zhang and Jin Si and Qiuli Mao and Shulin Zeng and Ke Hong and Genghan Zhang and Huazhong Yang and Yu Wang},
booktitle = {ASPLOS},
year = {2024},
doi = {10.1145/3620666.3651336}
}