一句话:在 Coq 中构建轻量框架,用一系列经验证的语义保持重写,对纯函数式数组语言写的张量核进行调度优化。
该工作针对张量核优化中调度变换的正确性问题,提出了一个基于 Coq 的轻量验证框架。与以往面向带数组与嵌套循环的命令式代码的编译不同,这里所有优化重写都是纯函数式语言内部的源到源变换,并已证明保持语义。语言包含一组用于表达高层计算细节的核心构造,以及从核心构造派生、但会触发存储模式与顺序等底层决策的 reshape 算子。用户通过一系列经验证的语义保持重写来编写调度。作者展示该框架不仅能推导出 Halide 等现有先进语言的优化、生成性能相当的代码,还能调度出 Halide 无法表达的一类有用程序变换。
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.
@article{paperbot1544,
title = {Verified tensor-program optimization via high-level scheduling rewrites},
author = {Amanda Liu and Gilbert Louis Bernstein and Adam Chlipala and Jonathan Ragan-Kelley},
journal = {Proceedings of the ACM on Programming Languages},
volume = {6},
number = {POPL},
year = {2022},
doi = {10.1145/3498717}
}