[论文解读] Relay: A High-Level IR for Deep Learning.

Jared Roesch, Steven Lyubomirsky+7 more|arXiv (Cornell University)|Apr 17, 2019

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

一句话总结

Relay 是一种用于深度学习的高级、函数式、静态类型的中间表示（IR），它统一并泛化了现有的深度学习 IR，实现了表达性强、可组合且可移植的模型编译。它通过支持自动微分和类型安全的优化、可扩展的编译，实现了在 CPU、GPU 和 FPGA 上的高性能，适用于最先进模型的编译。

ABSTRACT

Frameworks for writing, compiling, and optimizing deep learning (DL) models have recently enabled progress in areas like computer vision and natural language processing. Extending these frameworks to accommodate the rapidly diversifying landscape of DL models and hardware platforms presents challenging tradeoffs between expressiveness, composability, and portability. We present Relay, a new intermediate representation (IR) and compiler framework for DL models. The functional, statically-typed Relay IR unifies and generalizes existing DL IRs and can express state-of-the-art models. Relay's expressive IR required careful design of the type system, automatic differentiation, and optimizations. Relay's extensible compiler can eliminate abstraction overhead and target new hardware platforms. The design insights from Relay can be applied to existing frameworks to develop IRs that support extension without compromising on expressivity, composibility, and portability. Our evaluation demonstrates that the Relay prototype can already provide competitive performance for a broad class of models running on CPUs, GPUs, and FPGAs.

研究动机与目标

• 通过设计一种统一、可扩展的中间表示（IR），应对日益复杂和多样的深度学习模型与硬件平台。
• 在不牺牲性能或可维护性的前提下，平衡深度学习 IR 的表达性、可组合性和可移植性。
• 通过合理的 IR 设计，实现对 CPU、GPU 和 FPGA 上最先进模型的高效编译。

提出的方法

• 设计一种函数式、静态类型的 IR，泛化现有深度学习 IR，并支持复杂的模型结构。
• 将自动微分集成到 IR 中，原生支持反向传播和模型训练。
• 设计一种类型系统，确保正确性，并支持在异构模型组件之间进行安全优化。
• 构建一个可扩展的编译器流水线，消除抽象开销，并高效支持新硬件平台。
• 在类型化的函数式 IR 中应用优化技术，如公共子表达式消除和常量折叠。
• 通过平台特定的代码生成，实现 CPU、GPU 和 FPGA 之间模型的无缝编译与部署。

实验结果

研究问题

• RQ1 如何设计一种深度学习 IR，使其在统一并泛化现有表示的同时，保持表达性和可移植性？
• RQ2 需要哪些类型系统和语言特性，才能同时支持高级模型组合与高效的底层代码生成？
• RQ3 函数式、静态类型的 IR 是否能够在 CPU、GPU 和 FPGA 等多样化硬件平台上实现具有竞争力的性能？
• RQ4 如何在不牺牲优化机会的前提下，将自动微分原生集成到高级 IR 中？
• RQ5 在编译时能在多大程度上消除 IR 中的抽象，以实现与手写优化代码相当的性能？

主要发现

• Relay IR 能够成功表达最先进深度学习模型，包括在先前 IR 中难以表示的复杂架构。
• Relay 的函数式和静态类型特性提供了强大的正确性保证，并支持安全、可组合的优化。
• 原型编译器在 CPU、GPU 和 FPGA 上均实现了具有竞争力的性能，证明了其在多样化硬件平台上的可移植性。
• 该设计实现了高效的编译，抽象开销极小，性能接近手写优化实现。
• 可扩展的编译器架构使得新硬件后端和优化模块的集成变得简单。