Verified tensor-program optimization via high-level scheduling rewrites

We present a lightweight Coq framework for optimizing tensor kernels written in a pure, functional array language. Optimizations rely on user scheduling using series of verified, semantics-preserving rewrites. Unusually for compilation targeting imperative code with arrays and nested loops, all rewr...

Full description

Saved in:
Bibliographic Details
Published in:Proceedings of ACM on programming languages Vol. 6; no. POPL; pp. 1 - 28
Main Authors: Liu, Amanda, Bernstein, Gilbert Louis, Chlipala, Adam, Ragan-Kelley, Jonathan
Format: Journal Article
Language:English
Published: 01.01.2022
ISSN:2475-1421, 2475-1421
Online Access:Get full text
Tags: Add Tag
No Tags, Be the first to tag this record!
Description
Summary:We present a lightweight Coq framework for optimizing tensor kernels written in a pure, functional array language. Optimizations rely on user scheduling using series of verified, semantics-preserving rewrites. Unusually for compilation targeting imperative code with arrays and nested loops, all rewrites are source-to-source within a purely functional language. Our language comprises a set of core constructs for expressing high-level computation detail and a set of what we call reshape operators, which can be derived from core constructs but trigger low-level decisions about storage patterns and ordering. We demonstrate that not only is this system capable of deriving the optimizations of existing state-of-the-art languages like Halide and generating comparably performant code, it is also able to schedule a family of useful program transformations beyond what is reachable in Halide.
ISSN:2475-1421
2475-1421
DOI:10.1145/3498717