[论文解读] Sketch-Guided Equality Saturation: Scaling Equality Saturation to Complex Optimizations of Functional Programs

Thomas Kœhler, Phil Trinder+1 more|arXiv (Cornell University)|Nov 25, 2021

Parallel Computing and Optimization Techniques被引用 5

一句话总结

本文提出了一种基于草图引导的等式饱和技术，该技术结合高效的lambda演算编码与程序员提供的程序草图，以指导函数式程序优化中的等式饱和。该方法可在数秒内、内存占用低于1GB的情况下发现复杂的矩阵乘法优化（如循环分块、向量化和多线程化），而未经引导的等式饱和在运行一小时并消耗60GB内存后仍无法完成；

ABSTRACT

Generating high-performance code for diverse hardware and application domains is challenging. Functional array programming languages with patterns like map and reduce have been successfully combined with term rewriting to define and explore optimization spaces. However, deciding what sequence of rewrites to apply is hard and has a huge impact on the performance of the rewritten program. Equality saturation avoids the issue by exploring many possible ways to apply rewrites, efficiently representing many equivalent programs in an e-graph data structure. Equality saturation has some limitations when rewriting functional language terms, as currently naive encodings of the lambda calculus are used. We present new techniques for encoding polymorphically typed lambda calculi, and show that the efficient encoding reduces the runtime and memory consumption of equality saturation by orders of magnitude. Moreover, equality saturation does not yet scale to complex compiler optimizations. These emerge from long rewrite sequences of thousands of rewrite steps, and may use pathological combinations of rewrite rules that cause the e-graph to quickly grow too large. This paper introduces \emph{sketch-guided equality saturation}, a semi-automatic technique that allows programmers to provide program sketches to guide rewriting. Sketch-guided equality saturation is evaluated for seven complex matrix multiplication optimizations, including loop blocking, vectorization, and multi-threading. Even with efficient lambda calculus encoding, unguided equality saturation can locate only the two simplest of these optimizations, the remaining five are undiscovered even with an hour of compilation time and 60GB of RAM. By specifying three or fewer sketch guides all seven optimizations are found in seconds of compilation time, using under 1GB of RAM, and generating high performance code.

研究动机与目标

• 解决等式饱和在优化复杂函数式程序时的可扩展性限制，特别是那些需要数千步重写规则的场景。
• 克服由于结合了结合律和交换律等病态重写规则组合而导致的e-图指数级膨胀问题。
• 使等式饱和技术能够实际应用于如Rise这类具有多态类型的函数式语言，这类语言常用于高性能计算领域。
• 通过引入半自动的、基于草图的引导机制，减少重写策略中对手动阶段排序的依赖。
• 证明高效编码lambda演算对于使等式饱和在真实世界函数式程序优化中可行至关重要。

提出的方法

• 提出基于草图引导的等式饱和，程序员通过提供一系列程序草图（即带有未指定细节的程序模式）来引导优化过程。
• 将复杂优化分解为一系列较小的、可管理的等式饱和搜索，每个搜索均通过草图约束搜索空间。
• 开发一种使用De Bruijn索引和基于提取的替换机制的高效多态类型lambda演算编码，以最小化e-图膨胀并提升性能。
• 将新编码集成到Risegg等式饱和框架中，以支持带有名称绑定的函数语言项重写。
• 使用成本模型从e-图饱和后提取性能最高的程序，确保语义等价性和性能优化。
• 将该技术应用于Rise函数式语言中的七个矩阵乘法优化，与无引导等式饱和及TVM进行性能和资源使用对比评估。

实验结果

研究问题

• RQ1 基于草图引导的等式饱和能否有效将等式饱和扩展到原本无引导方法无法处理的复杂函数式程序优化？
• RQ2 高效的lambda演算编码在多大程度上减少了函数式程序优化中等式饱和的内存和运行时开销？
• RQ3 程序员提供的草图在多大程度上能引导发现复杂优化，而无需手动排序重写规则的阶段？
• RQ4 与无引导等式饱和及最先进的编译器TVM相比，基于草图引导的等式饱和在性能和资源使用方面带来了哪些改进？
• RQ5 高效编码与草图引导的结合是否能够使等式饱和发现原本因e-图爆炸而无法触及的优化？

主要发现

• 基于草图引导的等式饱和在低于1GB内存和数秒编译时间内，成功发现了全部七个复杂的矩阵乘法优化（包括循环分块、向量化和多线程化）。
• 无引导等式饱和即使在运行一小时并消耗60GB内存后，仍未能发现其中五个优化，凸显了朴素等式饱和的可扩展性局限。
• 新开发的lambda演算编码相比朴素编码将运行时间和内存消耗降低了数个数量级，使等式饱和在函数式程序中变得可行。
• 基于草图引导等式饱和生成的代码性能与最先进的TVM编译器相当，也与通过重写策略手动优化的代码性能相当。
• 仅使用三组或更少的草图引导，该技术即成功定位了所有复杂优化，证明了基于草图的引导在有效约束搜索空间方面的高效性。
• 该方法支持将声明式草图与等式饱和相结合，为函数式语言的交互式高性能代码生成提供了一条实用路径。