SparseAuto: An Auto-Scheduler for Sparse Tensor Computations Using Recursive Loop Nest Restructuring
Adhitha Dias, Logan Anderson, Kirshanthan Sundararajah, Artem, Pelenitsyn, Milind Kulkarni

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)…
Peer Reviews
No public reviews on file for this paper yet. If you reviewed it on a platform where reviews are public (OpenReview, ICLR, NeurIPS, ICML), you can paste yours below so the community can read it here.
Videos
No videos yet. Explain this paper in a talk, walkthrough, or lecture? Add one.
Taxonomy
TopicsParallel Computing and Optimization Techniques · Tensor decomposition and applications · Computational Physics and Python Applications
