[论文解读] Nimble: Efficiently Compiling Dynamic Neural Networks for Model Inference

Haichen Shen, Jared Roesch+7 more|arXiv (Cornell University)|Jun 4, 2020

Advanced Neural Network Applications参考文献 43被引用 25

一句话总结

Nimble 是一种新型深度学习编译器，通过引入动态类型系统、形状感知优化和轻量级虚拟机运行时，实现了对动态神经网络的高性能、可移植推理，支持控制流、动态张量形状和可变数据结构。它在 CPU、GPU 和 ARM 平台上相比现有解决方案最高可实现 20 倍的性能提升。

ABSTRACT

Modern deep neural networks increasingly make use of features such as dynamic control flow, data structures and dynamic tensor shapes. Existing deep learning systems focus on optimizing and executing static neural networks which assume a pre-determined model architecture and input data shapes--assumptions which are violated by dynamic neural networks. Therefore, executing dynamic models with deep learning systems is currently both inflexible and sub-optimal, if not impossible. Optimizing dynamic neural networks is more challenging than static neural networks; optimizations must consider all possible execution paths and tensor shapes. This paper proposes Nimble, a high-performance and flexible system to optimize, compile, and execute dynamic neural networks on multiple platforms. Nimble handles model dynamism by introducing a dynamic type system, a set of dynamism-oriented optimizations, and a light-weight virtual machine runtime. Our evaluation demonstrates that Nimble outperforms state-of-the-art deep learning frameworks and runtime systems for dynamic neural networks by up to 20x on hardware platforms including Intel CPUs, ARM CPUs, and Nvidia GPUs.

研究动机与目标

• 解决当前缺乏高效、可移植的动态神经网络推理系统的问题，这些系统需支持控制流、动态张量形状和可变数据结构。
• 克服现有深度学习框架和编译器的局限性，这些系统针对静态模型进行了优化，但在处理动态性方面表现不佳。
• 实现跨多种硬件平台（包括 Intel CPU、ARM CPU 和 NVIDIA GPU）的端到端动态模型优化与编译。
• 设计一个支持全模型优化、动态内存规划和可变张量形状及执行路径高效内核分发的系统。
• 提供一个可移植、轻量级的运行时，将平台无关的逻辑与硬件特定的内核解耦，确保跨平台兼容性和性能。

提出的方法

• 引入支持 'Any' 维度的动态类型系统，以实现对形状未知或可变的张量的类型推断与检查。
• 开发形状感知的代码生成技术，在编译时为动态变化的张量形状生成优化内核。
• 实现动态内存规划优化，静态分析并最小化所有可能执行路径中的运行时内存分配。
• 设计轻量级虚拟机（VM）运行时，将控制逻辑与硬件特定内核分离，实现跨平台的可移植性和高效分发。
• 集成设备放置启发式策略，根据张量形状和计算特性将操作分配到最优硬件加速器。
• 借鉴 LLVM 和 JVM 的编译技术，支持在静态可分析的中间表示（IR）中实现动态控制流和可变长度数据结构。

实验结果

研究问题

• RQ1 深度学习编译器如何有效表示和优化控制流、可变张量形状等动态神经网络结构？
• RQ2 在必须考虑所有可能执行路径和张量形状的前提下，为实现高性能，需要哪些优化？
• RQ3 轻量级、可移植的基于虚拟机的运行时是否能高效地在包括 CPU 和 GPU 在内的多种硬件平台上执行编译后的动态模型？
• RQ4 与静态或框架级方法相比，动态专用优化（如内存规划和形状感知代码生成）在多大程度上能提升推理性能？
• RQ5 全栈动态模型编译器在多种硬件平台和真实动态模型上的性能与现有解决方案相比如何？

主要发现

• 在主流硬件平台（包括 Intel CPU、ARM CPU 和 NVIDIA GPU）上，Nimble 相比最佳现有解决方案，延迟最低可降低 19.9 倍。
• 该系统在多种动态模型（包括 LSTM、Tree-LSTM 和 BERT）上均表现出一致的性能提升，加速比范围为 1.05 倍至 19.9 倍。
• 结合完整分发机制的动态形状感知代码生成，性能几乎与静态代码生成相当，同时避免了内核数量减少导致的性能下降。
• 在第二梯队硬件平台（如 ARM CPU）上，Nimble 超越了基于框架的解决方案，这些平台缺乏 cuDNN 或 oneDNN 等第三方库的可用或优化支持。
• 基于虚拟机的运行时实现了跨平台的可移植执行，且不损失性能，支持针对新张量形状和设备的按需代码生成。
• 动态内存规划通过在编译时分析所有可能的执行路径，显著减少了运行时内存分配开销，带来了可测量的性能提升。