A Verified Compiler for a Functional Tensor Language

Amanda Liu, Gilbert Bernstein, Adam Chlipala, Jonathan Ragan-Kelley. A Verified Compiler for a Functional Tensor Language. Proceedings of the 45th ACM SIGPLAN Conference on Programming Language Design and Implementation (PLDI'24). June 2024.

Paper as PDF


Producing efficient array code is crucial in high-performance domains like image processing and machine learning. This goal depends on an ability to control factors like compute intensity and locality by reordering computations into different stages and granularities with respect to where they are stored. However, traditional pure, functional tensor languages struggle to do so. In a previous publication, the ATL language was introduced as a pure, functional tensor language capable of systematically decoupling compute and storage order via a set of high-level combinators known as reshape operators. Reshape operators are a unique functional-programming construct since they manipulate storage location in the generated code by modifying the indices that appear on the left-hand sides of storage expressions. We present a formal correctness proof for an implementation of the compilation algorithm, marking the first verification of a lowering algorithm from a functional language that enables separate control of compute and storage ordering. One of the core difficulties of this proof required properly formulating the complex invariants to ensure that these storage-index remappings were well-formed. Notably, this revealed a soundness bug in the original published compilation algorithm regarding the truncation reshape operator. Our fix is a new type system that captures safety conditions that were previously implicit and enables us to prove compiler correctness for well-typed source programs. We evaluate this type system and compiler implementation on a range of common programs and optimizations, including but not limited to those previously studied, which demonstrated performance comparable to established compilers like Halide.

GitHub repository