[论文解读] Cortex: A Compiler for Recursive Deep Learning Models

Pratik Fegade, Tianqi Chen+2 more|arXiv (Cornell University)|Nov 2, 2020

Advanced Neural Network Applications参考文献 45被引用 10

一句话总结

Cortex 是一种编译器，通过启用端到端优化（如内核融合和模型持久化），将递归深度学习模型的推理延迟降至最低，从而在不依赖厂商库的情况下，在 CPU 和 GPU 上实现比以往工作低至 14 倍的推理延迟。

ABSTRACT

Optimizing deep learning models is generally performed in two steps: (i) high-level graph optimizations such as kernel fusion and (ii) low level kernel optimizations such as those found in vendor libraries. This approach often leaves significant performance on the table, especially for the case of recursive deep learning models. In this paper, we present Cortex, a compiler-based approach to generate highly-efficient code for recursive models for low latency inference. Our compiler approach and low reliance on vendor libraries enables us to perform end-to-end optimizations, leading to up to 14X lower inference latencies over past work, across different backends.

研究动机与目标

• 解决递归和动态深度学习模型在推理延迟方面的性能差距，尤其是在 CPU 和 GPU 上。
• 克服 cuDNN 和 MKL 等厂商库的局限性，这些库不支持 TreeLSTM 和 MV-RNN 等较少见的模型。
• 通过避免黑箱内核抽象，实现诸如内核融合和模型持久化等激进优化。
• 通过静态编译优化，支持递归模型在动态控制流下的高效执行。
• 提供一种可移植的、后端无关的编译栈，实现在各种硬件基础架构上的高性能。

提出的方法

• 通过观察控制流仅依赖于输入数据结构，将递归控制流降级为高效的基于循环的表示形式。
• 使用调度原语在张量计算之前执行静态优化，包括特化和动态批处理。
• 应用计算提升技术，将重复计算移出递归循环之外，减少冗余工作。
• 采用一种统一递归控制流和张量代数的中间表示（IR），实现端到端优化。
• 借鉴稀疏多面体框架和检查器-执行器模型的技术，表示并优化间接内存访问和数据结构布局。
• 将模型持久化和内核融合等优化形式化为编译器流水线中的头等转换原语。

实验结果

研究问题

• RQ1 基于编译器的方法是否能在优化递归深度学习模型的推理方面超越基于厂商库的框架？
• RQ2 如何在不引入运行时开销的情况下，在编译器 IR 中有效表示和优化递归控制流？
• RQ3 在避免使用厂商库抽象的前提下，端到端优化（如内核融合和模型持久化）在递归模型上可应用到何种程度？
• RQ4 静态优化（如动态批处理和计算提升）是否能在编译时有效应用于递归模型？
• RQ5 与现有框架相比，自定义编译器栈在不同硬件后端上的推理延迟性能表现如何？

主要发现

• Cortex 在包括 CPU 和 GPU 在内的多个后端上，实现的推理延迟相比以往工作最高降低 14 倍。
• 通过避免使用厂商库，Cortex 实现了激进的内核融合，减少了内核启动开销并改善了内存访问模式。
• 模型持久化被有效应用于递归模型，将频繁重用的参数保留在高速片上内存中，减少了对片外内存的访问。
• 基于编译器的方法支持静态动态批处理和特化，消除了 DyNet 和 PyTorch 等框架中常见的运行时开销。
• IR 设计成功地将递归控制流和张量计算统一表示，实现了两个领域之间的统一优化。
• Cortex 的优化具有可扩展性，可作为高层框架的底层后端，从而提升其在动态模型上的性能。