Verified tensor-program optimization via high-level scheduling rewrites
Name
3498717.pdf
Description
Published version
Size
344.58 KB
Format
Adobe PDF
Checksum (MD5)
2675f54caae7b758d21b56d2d81c0c75
Author(s) • • •
Liu, Amanda
Bernstein, Gilbert Louis
Chlipala, Adam
Ragan-Kelley, Jonathan
Date Issued
2022
Journal
Proceedings of the ACM on Programming Languages
Publisher
Association for Computing Machinery (ACM)
Citation
Liu, Amanda, Bernstein, Gilbert Louis, Chlipala, Adam and Ragan-Kelley, Jonathan. 2022. "Verified tensor-program optimization via high-level scheduling rewrites." Proceedings of the ACM on Programming Languages, 6 (POPL).
Version
Final published version
Abstract
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.
MIT Department
Massachusetts Institute of Technology. Computer Science and Artificial Intelligence Laboratory
Terms of Use
Creative Commons Attribution 4.0 International license
Persistent DSpace Link
DOI of Published Version
https://doi.org/10.1145/3498717