We address the problem of optimizing sparse tensor algebra in a compiler and show how to define standard loop transformations (such as splitting, collapsing, and reordering) for sparse iteration spaces. The key idea is that these transformations map sparse iteration spaces into \textit{derived} iteration spaces. This results in a new iteration order that respects the constraints inherent in sparse data structures. We further demonstrate that derived iteration spaces can be used to tile over both the indexes (coordinates) and the nonzeros; the former is equivalent to tiling dense iteration spaces, while the latter tiles sparse iteration spaces into statically load-balanced blocks of nonzeros. Tiling the space of nonzeros lets the compiler efficiently exploit heterogeneous compute resources such as threads, vector units, and GPUs.

We implemented these concepts and an associated scheduling API on top of the open-source TACO system. The scheduling API can be used manually by performance engineers or be the target of an automatic scheduling system. We outline a heuristics-based autoscheduling system in this paper, but many other systems are possible and enabled by our API. Using the scheduling API, we show how to optimize sparse and mixed sparse-dense tensor algebra expressions on both CPUs and GPUs. Our results show that the sparse transformations are sufficient to generate code with competitive performance to hand-optimized implementations from the literature, while generalizing to all of the tensor algebra.