[论文解读] SPORES: Sum-Product Optimization via Relational Equality Saturation for Large Scale Linear Algebra

Yisu Remy Wang, Shana Hutchison+3 more|arXiv (Cornell University)|Feb 19, 2020

Parallel Computing and Optimization Techniques参考文献 27被引用 11

一句话总结

SPORES 通过等式饱和和 E-Graphs 提出了一套基于关系代数的完整优化框架，用于大规模线性代数工作负载，系统性地探索等价表达式。它集成到 SystemML 中，重新发现所有现有优化并发现新优化，在真实机器学习工作负载中实现了高达 5 倍的性能提升，充分利用了稀疏性和代数等价性。

ABSTRACT

Machine learning algorithms are commonly specified in linear algebra (LA). LA expressions can be rewritten into more efficient forms, by taking advantage of input properties such as sparsity, as well as program properties such as common subexpressions and fusible operators. The complex interaction among these properties' impact on the execution cost poses a challenge to optimizing compilers. Existing compilers resort to intricate heuristics that complicate the codebase and add maintenance cost but fail to search through the large space of equivalent LA expressions to find the cheapest one. We introduce a general optimization technique for LA expressions, by converting the LA expressions into Relational Algebra (RA) expressions, optimizing the latter, then converting the result back to (optimized) LA expressions. One major advantage of this method is that it is complete, meaning that any equivalent LA expression can be found using the equivalence rules in RA. The challenge is the major size of the search space, and we address this by adopting and extending a technique used in compilers, called equality saturation. We integrate the optimizer into SystemML and validate it empirically across a spectrum of machine learning tasks; we show that we can derive all existing hand-coded optimizations in SystemML, and perform new optimizations that lead to speedups from 1.2X to 5X.

研究动机与目标

• 解决基于启发式规则的编译器在优化复杂线性代数表达式时，于稀疏性、公共子表达式和可融合算子等多维优化空间中的局限性。
• 通过将关系代数作为中间表示，实现对等价优化线性代数表达式的完整、系统性搜索。
• 通过等式饱和和 E-Graphs 管理等价表达式搜索空间的不可行规模，实现在不牺牲完整性的前提下可扩展的探索。
• 将优化器集成到 SystemML 中，并通过实证验证其重新发现已知优化和发现可提升性能的新优化的能力。

提出的方法

• 使用正式映射将线性代数表达式转换为关系代数表达式，将求和、乘积和转置操作表示为关系运算符（例如连接、并集、聚集）。
• 在 K-关系（具有数值属性的多重集语义）上使用关系代数框架，通过连接查询的同构性确保等价性可判定且完整性可实现。
• 应用等式饱和技术，利用一组标准关系恒等式填充 E-Graph，紧凑地表示搜索空间中所有等价表达式。
• 利用基于模式的类不变量访问属性信息，实现在规则应用过程中进行常量折叠和精确成本估算。
• 在等式饱和过程中采样重写匹配，以平衡探索深度与收敛时间，避免在大规模程序中进行完全饱和。
• 使用基于成本的约束求解器从 E-Graph 中提取最优表达式，然后将其转换回线性代数表达式以供执行。

实验结果

研究问题

• RQ1 基于关系代数的中间表示是否能够实现对等价线性代数表达式的完整探索，包括启发式规则所遗漏的表达式？
• RQ2 如何将等式饱和和 E-Graphs 适配以管理由并集、聚集和公共子表达式引入的和-积表达式中的大规模搜索空间？
• RQ3 该方法在多大程度上能够重新发现并超越生产级编译器（如 SystemML）中手写优化的性能？
• RQ4 通过在真实世界机器学习工作负载中系统性搜索，可实现多大程度的性能提升，发现新的非平凡优化？

主要发现

• SPORES 成功重新发现 SystemML 中所有现有手写优化，证明了其完整性和实际适用性。
• 优化器发现了此前未知的新优化，使多种机器学习工作负载的性能提升范围达到 1.2× 至 5×。
• 通过使用具有多重集语义和规范形式等价性的关系代数，SPORES 确保所有等价表达式均被系统性探索，实现了理论上的完整性。
• 类不变量的集成实现了精确的成本估算和有效的常量折叠，提升了优化决策的质量。
• 在等式饱和过程中采样重写匹配，在不进行完全饱和的前提下实现了良好的探索效率，使该方法可扩展至大型程序。
• 实证评估表明，尽管约束求解器在理论上最优，但引入了显著开销；因此，基于采样的探索在实践中更为高效。