[论文解读] Bring Your Own Codegen to Deep Learning Compiler

Zhi Chen, Cody Hao Yu+7 more|arXiv (Cornell University)|May 3, 2021

Advanced Neural Network Applications参考文献 31被引用 7

一句话总结

本文提出了一种统一的开源框架，使深度学习加速器厂商能够将其专有的代码生成工具集成到现有的深度学习编译器中，从而将特定于加速器的代码生成与共享的编译器基础设施解耦。通过重用标准编译器优化并支持灵活的模型划分，该框架将开发工作量减少至仅数千行代码，并已成功部署于NVIDIA Jetson和Xilinx Vitis-AI等商业系统中，实现了可测量的性能提升。

ABSTRACT

Deep neural networks (DNNs) have been ubiquitously applied in many applications, and accelerators are emerged as an enabler to support the fast and efficient inference tasks of these applications. However, to achieve high model coverage with high performance, each accelerator vendor has to develop a full compiler stack to ingest, optimize, and execute the DNNs. This poses significant challenges in the development and maintenance of the software stack. In addition, the vendors have to contiguously update their hardware and/or software to cope with the rapid evolution of the DNN model architectures and operators. To address these issues, this paper proposes an open source framework that enables users to only concentrate on the development of their proprietary code generation tools by reusing as many as possible components in the existing deep learning compilers. Our framework provides users flexible and easy-to-use interfaces to partition their models into segments that can be executed on "the best" processors to take advantage of the powerful computation capability of accelerators. Our case study shows that our framework has been deployed in multiple commercial vendors' compiler stacks with only a few thousand lines of code.

研究动机与目标

• 解决为DNN加速器开发完整编译器栈的高昂工程成本问题，尽管各厂商存在共享的优化需求，但该成本仍被重复投入。
• 使加速器厂商能够专注于其专有的代码生成，同时重用标准编译器优化和模型划分技术。
• 解决模型不兼容、复杂算子执行效率低下以及各厂商重复实现通用编译器组件的挑战。
• 提供一种灵活且可扩展的框架，支持多种加速器，并与现有深度学习编译器生态系统无缝集成。
• 通过抽象平台无关的优化并支持加速器特定代码生成的模块化集成，减少部署时间和维护开销。

提出的方法

• 将深度学习编译器分解为两部分：共享的、平台无关的流水线和仅针对加速器的代码生成模块。
• 重用现有的深度学习编译器后端（如ONNX Runtime、Glow）来处理模型加载、图优化和运行时执行。
• 引入一种模型划分机制，根据算子支持情况和性能特征将DNN划分为子图。
• 为厂商提供标准化接口，用于插入其自有的代码生成逻辑，从而抽象底层硬件细节。
• 在共享流水线中支持硬件感知优化，如布局转换、量化和公共子表达式消除。
• 通过将生成的代码与轻量级运行时集成，实现端到端的编译和执行，该运行时负责数据移动和内核启动管理。

实验结果

研究问题

• RQ1 加速器厂商如何降低为DNN推理构建完整编译器栈的工程负担？
• RQ2 在支持厂商特定代码生成的同时，共享编译器优化在多种加速器上可重用的程度如何？
• RQ3 模块化框架能否实现对包含控制流和非线性算子的复杂DNN模型的高效且高性能执行？
• RQ4 该框架如何简化NVIDIA Jetson和Xilinx Vitis-AI等商业加速器的集成？
• RQ5 与从零开始构建自定义编译器栈相比，该框架在性能和开发工作量方面带来了哪些改进？

主要发现

• 该框架使商业加速器厂商仅需数千行代码即可将其代码生成模块集成到最先进的编译器中。
• 多家厂商已成功将该框架部署于生产级编译器栈中，包括NVIDIA Jetson Xavier和Xilinx Vitis-AI。
• 通过利用共享优化和高效的模型划分，该框架在性能上显著优于基线实现。
• 通过将代码生成与共享基础设施解耦，该框架降低了维护成本，并加快了新加速器的上市速度。
• 该方法通过将复杂模型中的控制流和非线性算子卸载至CPU或运行时执行，同时在加速器上加速计算密集型内核，从而支持复杂模型的执行。
• 该框架在包括FPGA、ASIC和通用加速器在内的多种硬件目标上均表现出良好的兼容性和可重用性。