A High Performance Sparse Tensor Algebra Compiler in MLIR
Parallel Programming Systems
TimeSunday, 14 November 20213:30pm - 4:10pm CST
DescriptionSparse tensor algebra is widely used in many applications. The performance of sparse tensor algebra kernels strongly depends on the characteristics of the input tensors, hence many storage formats are designed for tensors to achieve optimal performance for particular applications/architectures, which makes it challenging to implement and optimize every tensor operation of interest on a given architecture. We propose a tensor algebra domain-specific language (DSL) and compiler framework to automatically generate kernels for mixed sparse-dense tensor algebra operations.
The proposed DSL provides high-level programming abstractions to represent tensor algebra operations. The compiler introduces a new Sparse Tensor Algebra dialect built on top of LLVM's extensible MLIR compiler infrastructure for efficient code generation while covering a wide range of tensor formats. It also leverages input-dependent code optimization to enhance data locality for better performance. Our results show that the performance of automatically generated kernels outperforms the state-of-the-art sparse tensor algebra compiler.