[论文解读] Tenspiler: A Verified Lifting-Based Compiler for Tensor Operations (Extended Version)

Jie Qiu, C.R. Cai+4 more|arXiv (Cornell University)|Jan 1, 2024

Parallel Computing and Optimization Techniques被引用 3

一句话总结

Tenspiler 是一个经过验证的基于提升的编译器，通过一种新颖的统一中间语言 Tensir，将顺序的 C++ 或 Python 代码转换为优化的张量操作。通过将代码转换问题形式化为语法引导的程序综合问题，并利用定理证明确保正确性，Tenspiler 在六个软件和硬件后端的 69 个真实世界基准上实现了平均 105× 的内核加速和 9.65× 的端到端加速。

ABSTRACT

Tensor processing infrastructures such as deep learning frameworks and specialized hardware accelerators have revolutionized how computationally intensive code from domains such as deep learning and image processing is executed and optimized. These infrastructures provide powerful and expressive abstractions while ensuring high performance. However, to utilize them, code must be written specifically using the APIs / ISAs of such software frameworks or hardware accelerators. Importantly, given the fast pace of innovation in these domains, code written today quickly becomes legacy as new frameworks and accelerators are developed, and migrating such legacy code manually is a considerable effort. To enable developers in leveraging such DSLs while preserving their current programming paradigm, we introduce Tenspiler, a verified lifting-based compiler that uses program synthesis to translate sequential programs written in general-purpose programming languages (e.g., C++ or Python code) into tensor operations. Central to Tenspiler is our carefully crafted yet simple intermediate language, named TensIR, that expresses tensor operations. TensIR enables efficient lifting, verification, and code generation. Currently, Tenspiler already supports $ extbf{six}$ DSLs, spanning a broad spectrum of software and hardware environments. Furthermore, we show that new backends can be easily supported by Tenspiler by adding simple pattern-matching rules for TensIR. Using 10 real-world code benchmark suites, our experimental evaluation shows that by translating code to be executed on $ extbf{6}$ different software frameworks and hardware devices, Tenspiler offers on average 105$ imes$ kernel and 9.65$ imes$ end-to-end execution time improvement over the fully-optimized sequential implementation of the same benchmarks.

研究动机与目标

• 通过实现从通用代码到领域特定语言（DSL）目标的自动转换，解决张量计算中遗留代码的可移植性挑战。
• 克服基于模式匹配的编译器和基于大语言模型（LLM）的代码生成的局限性，后者脆弱、易出错且缺乏正确性保证。
• 设计一种统一且可扩展的中间表示（Tensir），抽象张量操作，以实现高效合成与验证。
• 通过向 Tensir 添加简单的模式匹配规则，实现对新后端的无缝支持，从而显著降低移植到新框架或加速器的工作量。

提出的方法

• 将代码转换任务形式化为语法引导的程序合成（SGS）问题，将输入程序提升为等价的 Tensir 程序。
• 使用定理证明器验证输入程序与合成的 Tensir 程序之间的语义等价性，确保正确性。
• 设计 Tensir 为一种最小化、基于张量代数的中间表示，捕捉常见的向量和矩阵操作，从而实现高效合成与验证。
• 应用如操作符限制和程序状态限制等优化技术，以缩小搜索空间并提升可扩展性。
• 使用用户提供的代码生成器，从经过验证的 Tensir 程序生成目标代码，支持包括 MLX、Gemmini 和 Gaudi 在内的多个后端。
• 利用表达式树和向量化启发式方法，指导高维张量中高效向量化操作的合成。

实验结果

研究问题

• RQ1 基于验证的提升方法是否能够实现通用代码到多样化张量 DSL 和硬件后端的正确且高效的编译？
• RQ2 统一的中间表示（Tensir）在支持异构张量处理目标的可扩展合成与验证方面有多高效？
• RQ3 Tenspiler 的合成与验证流水线在涉及 1D、2D 和 3D 张量操作的真实世界复杂基准上，其可扩展性如何？
• RQ4 与完全优化的顺序实现相比，Tenspiler 在内核执行时间和端到端执行时间方面的性能表现如何？
• RQ5 Tenspiler 是否能够以极低的工程成本轻松扩展以支持新后端？

主要发现

• 在 69 个真实世界基准上，Tenspiler 相较于完全优化的顺序实现，实现了平均 105× 的内核执行时间加速和 9.65× 的端到端执行时间加速。
• 合成过程可有效扩展，深度为 5 的人工 3D 张量基准合成耗时 184 秒，远在 1 小时超时限制之内。
• Tenspiler 成功将代码转译到六个不同的后端，包括 MLX、Gemmini 和 Gaudi，展示了其与开源及商业框架的广泛兼容性。
• 为新发布的 MLX 框架实现 Tenspiler 所需代码少于 200 行，凸显了 Tensir 中间表示的可扩展性。
• LLM 生成的相同基准代码始终存在错误且无法运行，凸显了 LLM 在为不熟悉的 DSL 生成经验证代码方面的局限性。
• 使用 Tensir 能够实现高效的验证与代码生成，其中操作符限制和程序状态限制等优化手段显著缩小了合成的搜索空间。