SparseAuto: An Auto-Scheduler for Sparse Tensor Computations Using Recursive Loop Nest Restructuring

TL;DR

SparseAuto is an auto-scheduler that optimizes sparse tensor computations by using recursive loop restructuring, kernel fission and fusion, and SMT-based search to generate high-performance code with significant speedups.

Contribution

It introduces a novel auto-scheduler employing generalized loop transformations and SMT solvers to produce asymptotically better sparse tensor computation schedules.

Findings

Achieves orders-of-magnitude speedup over TACO-generated code

Uses SMT solvers to approximate Pareto-optimal schedules

Effectively reduces asymptotic complexity and memory footprint

Abstract

Automated code generation and performance enhancements for sparse tensor algebra have become essential in many real-world applications, such as quantum computing, physical simulations, computational chemistry, and machine learning. General sparse tensor algebra compilers are not always versatile enough to generate asymptotically optimal code for sparse tensor contractions. This paper shows how to generate asymptotically better schedules for complex sparse tensor expressions using kernel fission and fusion. We present generalized loop restructuring transformations to reduce asymptotic time complexity and memory footprint. Furthermore, we present an auto-scheduler that uses a partially ordered set (poset)-based cost model that uses both time and auxiliary memory complexities to prune the search space of schedules. In addition, we highlight the use of Satisfiability Module Theory (SMT) solvers in sparse auto-schedulers to approximate the Pareto frontier of better schedules to the smallest number of possible schedules, with user-defined constraints available at compile-time. Finally, we show that our auto-scheduler can select better-performing schedules and generate code for them. Our results show that the auto-scheduler provided schedules achieve orders-of-magnitude speedup compared to the code generated by the Tensor Algebra Compiler (TACO) for several computations on different real-world tensors.