[论文解读] Linnea: Automatic Generation of Efficient Linear Algebra Programs

Henrik Barthels, Christos Psarras+1 more|arXiv (Cornell University)|Dec 30, 2019

Parallel Computing and Optimization Techniques参考文献 50被引用 4

一句话总结

Linnea 是一种代码生成器，可将高层线性代数表达式自动转换为优化的 BLAS 和 LAPACK 内核调用序列。通过使用符号重写规则和最佳优先搜索算法，它生成高效且正确构造的代码，其性能在 125 个测试问题中比 MATLAB、Julia、Eigen 和 Armadillo 等高级语言高出最多 10 倍。

ABSTRACT

The translation of linear algebra computations into efficient sequences of library calls is a non-trivial task that requires expertise in both linear algebra and high-performance computing. Almost all high-level languages and libraries for matrix computations (e.g., Matlab, Eigen) internally use optimized kernels such as those provided by BLAS and LAPACK; however, their translation algorithms are often too simplistic and thus lead to a suboptimal use of said kernels, resulting in significant performance losses. In order to combine the productivity offered by high-level languages, and the performance of low-level kernels, we are developing Linnea, a code generator for linear algebra problems. As input, Linnea takes a high-level description of a linear algebra problem; as output, it returns an efficient sequence of calls to high-performance kernels. Linnea uses a custom best-first search algorithm to find a first solution in less than a second, and increasingly better solutions when given more time. In 125 test problems, the code generated by Linnea almost always outperforms Matlab, Julia, Eigen and Armadillo, with speedups up to and exceeding 10x.

研究动机与目标

• 解决高层线性代数抽象与低层优化内核（如 BLAS 和 LAPACK）之间的性能差距。
• 自动化将数学表达式转换为高效且正确的库调用序列，而无需高性能计算领域的专家知识。
• 减少手动优化线性代数代码以提升性能所带来的手动工作量和易错性。
• 通过结合符号优化和领域特定知识，使非专家和专家均能生成高性能代码。
• 通过结合高层表达力与低层性能，提升科学计算中的生产力。

提出的方法

• 接受高层线性代数表达式作为输入，包括矩阵、向量、标量以及加法、乘法、转置和求逆等操作。
• 使用符号操作和重写规则，应用代数恒等式（例如分配律、子表达式重用）以实现性能优化。
• 采用自定义的最佳优先搜索算法，探索可能的代码序列空间，优先选择高效解。
• 输出 Julia 代码，利用 Julia 对 BLAS/LAPACK 的高效封装和可扩展性以实现高性能。
• 支持矩阵属性（例如 SPD、三角形、对称）作为注解，以指导内核选择和优化。
• 使用成本模型引导搜索，并通过合并搜索分支来随时间提升解的质量。

实验结果

研究问题

• RQ1 符号代码生成器能否自动生成比 MATLAB、Julia、Eigen 和 Armadillo 等高级语言更高效的线性代数代码？
• RQ2 结合符号重写规则的最佳优先搜索策略在发现高性能代码序列方面的有效性如何？
• RQ3 自动优化线性代数表达式在多大程度上可以减少因简单翻译导致的性能差距？
• RQ4 引入矩阵属性（例如 SPD、三角形）在多大程度上能提升代码生成质量？
• RQ5 在自动线性代数代码合成的背景下，生成时间与代码质量之间的权衡是什么？

主要发现

• 在 125 个测试问题中，Linnea 生成的代码在性能上持续优于 MATLAB、Julia、Eigen 和 Armadillo，最高提速达 10 倍以上。
• Linnea 在不到一秒内找到解，并在获得更多时间后持续优化，表明其在优化空间中具有高效的搜索能力。
• 符号重写规则的使用实现了正确构造的代码生成，避免了手动优化中常见的正确性问题。
• 矩阵属性（例如 SPD、三角形）的集成使 Linnea 能够选择最优内核（如 Cholesky 分解或专用求解器），从而降低计算成本。
• 即使在 Eigen 和 Armadillo 等库使用优化函数的情况下，生成的代码仍比其朴素实现和推荐实现更高效。
• 成本模型有效引导了搜索过程，且通过合并搜索分支显著提升了随时间推移的解质量。