为AI编译器中的模式匹配提供形式化基础,让优化规则更可靠
Pattern Matching in AI Compilers and its Formalization (Extended Version)
- 用可递归、带约束的模式语言匹配计算图子结构
- 建立形式化语义与算法语义,并证明两者等价
- 适合编译器开发者和形式化验证研究者
PyPM 是一种基于 Python 的领域特定语言,用于在机器学习计算图上构建基于重写的优化模块。用户通过定义(a)匹配计算图子图的模式和(b)将匹配子图替换为优化内核的规则来实现优化。与众多其他重写语言不同,PyPM 的模式语言具有复杂且新颖的特性,借鉴了逻辑编程思想,支持递归、非确定性匹配,并可检查张量形状等特定领域约束。因此其实现极为复杂,包含数千行 C++ 代码。本文介绍 PyPM 的构建工作,以及对其复杂性的形式化与提炼,提出一个表达 PyPM 模式语言核心操作的形式化核心演算。我们定义了声明式语义(描述哪些模式匹配哪些项)和算法式语义(理想化的模式解释器版本),并证明二者等价。整个形式化过程在 Coq 证明助手内完全机械化完成。
原文摘要 · Abstract (English)
PyPM is a Python-based domain specific language (DSL) for building rewrite-based optimization passes on machine learning computation graphs. Users define individual optimizations by writing (a) patterns that match subgraphs of a computation graph and (b) corresponding rules which replace a matched subgraph with an optimized kernel. PyPM is distinguished from the many other DSLs for defining rewriting passes by its complex and novel pattern language which borrows concepts from logic programming. PyPM patterns can be recursive, nondeterminstic, and can require checking domain-specific constraints such as the shapes of tensors. The PyPM implementation is thus similarly complicated, consisting of thousands of lines of C++ code. In this paper, we present our work on building PyPM, as well as formalizing and distilling and this complexity to an understandable mathematical core. We have developed a formal core calculus expressing the main operations of the PyPM pattern language. We define both a declarative semantics - describing which patterns match which terms - and an algorithmic semantics - an idealized version of the PyPM pattern interpreter - and prove their equivalence. The development is fully mechanized in the Coq proof assistant.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。