An Asymptotic Cost Model for Autoscheduling Sparse Tensor Programs

TL;DR

This paper introduces the first automatic asymptotic scheduler for sparse tensor programs, significantly reducing scheduling complexity and improving performance by selecting optimal kernels based on asymptotic cost analysis.

Contribution

It presents a novel approach to automatically select asymptotically optimal schedules for sparse tensor programs, narrowing the search space to a Pareto frontier.

Findings

Reduces scheduling space by orders of magnitude

Generated kernels outperform default schedules asymptotically

Automates decision-making in sparse tensor compilation

Abstract

While loop reordering and fusion can make big impacts on the constant-factor performance of dense tensor programs, the effects on sparse tensor programs are asymptotic, often leading to orders of magnitude performance differences in practice. Sparse tensors also introduce a choice of compressed storage formats that can have asymptotic effects. Research into sparse tensor compilers has led to simplified languages that express these tradeoffs, but the user is expected to provide a schedule that makes the decisions. This is challenging because schedulers must anticipate the interaction between sparse formats, loop structure, potential sparsity patterns, and the compiler itself. Automating this decision making process stands to finally make sparse tensor compilers accessible to end users. We present, to the best of our knowledge, the first automatic asymptotic scheduler for sparse tensor programs. We provide an approach to abstractly represent the asymptotic cost of schedules and to choose between them. We narrow down the search space to a manageably small "Pareto frontier" of asymptotically undominated kernels. We test our approach by compiling these kernels with the TACO sparse tensor compiler and comparing them with those generated with the default TACO schedules. Our results show that our approach reduces the scheduling space by orders of magnitude and that the generated kernels perform asymptotically better than those generated using the default schedules.