[论文解读] A High-Performance Sparse Tensor Algebra Compiler in Multi-Level IR

Ruiqin Tian, Luanzheng Guo+3 more|arXiv (Cornell University)|Feb 9, 2021

Parallel Computing and Optimization Techniques参考文献 65被引用 10

一句话总结

COMET 是一个基于 MLIR 的高性能稀疏张量代数编译器，采用类似爱因斯坦记号的领域特定语言（DSL）抽象，自动为混合稀疏-密集张量操作生成优化内核。它通过四个维度属性支持多种存储格式，并利用数据重排技术提升数据局部性，在 SpMV 上相比 TACO 最快提升 20.92 倍，在 SpMM 上提升 6.39 倍，在 TTM 上提升 13.9 倍。

ABSTRACT

Tensor algebra is widely used in many applications, such as scientific computing, machine learning, and data analytics. The tensors represented real-world data are usually large and sparse. There are tens of storage formats designed for sparse matrices and/or tensors and the performance of sparse tensor operations depends on a particular architecture and/or selected sparse format, which makes it challenging to implement and optimize every tensor operation of interest and transfer the code from one architecture to another. We propose a tensor algebra domain-specific language (DSL) and compiler infrastructure to automatically generate kernels for mixed sparse-dense tensor algebra operations, named COMET. The proposed DSL provides high-level programming abstractions that resemble the familiar Einstein notation to represent tensor algebra operations. The compiler performs code optimizations and transformations for efficient code generation while covering a wide range of tensor storage formats. COMET compiler also leverages data reordering to improve spatial or temporal locality for better performance. Our results show that the performance of automatically generated kernels outperforms the state-of-the-art sparse tensor algebra compiler, with up to 20.92x, 6.39x, and 13.9x performance improvement, for parallel SpMV, SpMM, and TTM over TACO, respectively.

研究动机与目标

• 解决在不同存储格式和架构上手动优化稀疏张量内核的挑战。
• 通过类似爱因斯坦记号的 DSL 实现张量代数操作的高效率编程。
• 通过与 MLIR 框架集成，实现性能可移植性和可扩展性。
• 为广泛存储格式下的混合稀疏-密集张量操作提供高效代码生成。
• 通过编译管道中自动化的数据重排技术，提升数据局部性和性能。

提出的方法

• 设计一种领域特定语言（DSL），支持使用类似爱因斯坦记号的语法表达高层张量代数表达式。
• 基于 MLIR 的多级中间表示（IR），将领域特定优化与架构特定优化解耦。
• 使用四个维度属性表示张量存储格式：密集、压缩唯一、压缩非唯一和单例。
• 实现一种代码生成算法，通过分析维度属性，生成高效、格式感知的内核，用于非零元素的迭代。
• 集成一种数据重排算法，在内核生成过程中提升空间和时间局部性。
• 在 MLIR 层次结构的高层应用架构无关优化，在低层实施寄存器/内存优化。

实验结果

研究问题

• RQ1 如何设计一个张量代数编译器，以在保持高性能的同时支持广泛的稀疏张量存储格式？
• RQ2 基于爱因斯坦记号的 DSL 在稀疏张量计算中能在多大程度上提升程序员生产力？
• RQ3 基于 MLIR 的编译能否实现在多种架构和格式之间的性能可移植性与可扩展性？
• RQ4 自动化数据重排在提升稀疏张量操作的数据局部性和内核性能方面有多有效？
• RQ5 与现有编译器（如 TACO）相比，COMET 在关键稀疏张量内核（如 SpMV、SpMM 和 TTM）上能实现多大的性能提升？

主要发现

• COMET 在 CPU 上的稀疏矩阵-向量乘法（SpMV）中相比 TACO 最快提升 20.92 倍。
• COMET 在 CPU 上的稀疏张量-矩阵乘法（SpMM）中相比 TACO 提升 6.39 倍。
• COMET 在 CPU 上的张量变换（TTM）操作中相比 TACO 提升 13.9 倍。
• 数据重排的使用显著改善了数据局部性，是性能提升的关键因素。
• 通过四个维度属性实现的格式无关设计，使 COMET 能够无需硬编码每种组合，即支持 COO、CSR、DCSR、ELLPACK、CSF 等常见格式及通用模式。
• 基于 MLIR 的架构支持模块化、可扩展和可移植的编译，便于未来支持 GPU 等异构架构。