Sparse and Structured Tensor Programming

Abstract

From FORTRAN to NumPy, tensors have revolutionized how we express computation. However, tensors in these, and almost all prominent systems, can only handle dense rectilinear grids of values. Real-world tensors are often structured, containing patterns which allow us to optimize storage or computation, such as sparsity (mostly zero), runs of repeated values, or symmetry. Specializing implementations for structure yields significant speedups, but support for structured tensors is fragmented and incomplete. The heart of the problem is coiteration, simultaneously iterating over multiple tensors in a program, where each tensor format may have different internal structure. As each combination of structures requires a unique coiteration algorithm, existing frameworks struggle to abstract over the design space, instead hard-coding support for a few programs and/or a few structures. In this thesis, we build an abstraction for coiteration, enabling us to support both a wide range of programs and diverse tensor structures. We use a language, looplets, to describe the structure of tensors in tensor programs. Looplets allow the compiler to generate code to coiterate over any combination of structured tensor formats. The looplets language decomposes loops over sparse and structured formats hierarchically. This decomposition simplifies compilation, allowing us to capture key mathematical properties (such as x∗0 = 0, which motivates sparsity) with simple term rewriting. Building on looplets, we introduce a new language, Finch, for general structured tensor programming. Finch makes it easier to compute with structured tensors by combining program control flow and tensor structures into a common representation where they can be co-optimized. Finch automatically specializes control flow to data so that performance engineers can focus on experimenting with many algorithms. Finch supports a familiar programming language of loops, statements, ifs, breaks, etc., over a wide variety of tensor structures, such as sparsity, run-length-encoding, symmetry, triangles, padding, or blocks. Finch reliably utilizes the key properties of each structure, making it easier to write and optimize structured tensor programs. In our case studies, we show that this leads to dramatic speedups in diverse applications, including linear algebra, image processing, and graph analytics. Our abstracted design makes it easier to extend Finch to new tensor structures and programming models. Finch has been separately extended to support a DSL for symmetry-aware tensor programs and to support real-valued indexing.