[论文解读] Generalizing Shape Analysis with Gradual Types

Migeed, Zeina, Reed, James+2 more|arXiv (Cornell University)|Apr 17, 2019

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

一句话总结

Relay 是一种高层级、静态类型、函数式中间表示（IR），专为深度学习编译器设计，统一并泛化了现有的 IR，以支持表达性强的模型、可组合的优化以及可移植的硬件目标。它通过将领域特定优化重新定义为标准编译器阶段，并支持可扩展的、平台无关的代码生成，实现了在 CPU、GPU 和加速器上的竞争力性能。

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 expressivity, composability, and portability. We present Relay, a new compiler framework for DL. Relay's functional, statically typed intermediate representation (IR) unifies and generalizes existing DL IRs to express state-of-the-art models. The introduction of Relay's expressive IR requires careful design of domain-specific optimizations, addressed via Relay's extension mechanisms. Using these extension mechanisms, Relay supports a unified compiler that can target a variety of hardware platforms. Our evaluation demonstrates Relay's competitive performance for a broad class of models and devices (CPUs, GPUs, and emerging accelerators). Relay's design demonstrates how a unified IR can provide expressivity, composability, and portability without compromising performance.

研究动机与目标

• 解决在统一、可扩展的编译器框架中支持多样化深度学习模型和硬件平台的日益增长的挑战。
• 克服现有深度学习 IR 在扩展新模型、优化或硬件时牺牲表达性、可组合性或可移植性的局限。
• 在不损害性能的前提下，实现领域特定优化（如量化、融合）与硬件特定代码生成的无缝集成。
• 通过统一模型表示、优化与代码生成，为未来深度学习编译器研究提供一种原则性且可扩展的基础。

提出的方法

• 设计 Relay 为一种支持头等函数、控制流以及复杂数据结构（如树和图）的函数式、静态类型 IR。
• 利用成熟的编译技术，将常见的深度学习优化（如量化、形状推断、算子融合）重新定义为标准编译器阶段。
• 引入可扩展机制以定义新算子和优化，支持编译器栈的模块化与可组合性扩展。
• 实现平台无关的算子表示，使同一 IR 能够针对多种后端（包括 CPU、GPU 和定制加速器）进行目标生成。
• 使用统一的 IR 将包含复杂控制流（如 RNN、TreeLSTM）的完整模型编译为轻量级、高度优化的二进制文件，无需依赖宿主语言脚本。
• 利用数十年的传统编译器研究成果，以正确且高效的方式组合优化，避免实现中的组合爆炸问题。

实验结果

研究问题

• RQ1 是否可以通过函数式、静态类型的 IR 统一并泛化现有的深度学习 IR，以表达具有复杂控制流和数据结构的最先进模型？
• RQ2 当将领域特定优化（如量化和算子融合）重新定义为标准编译器阶段时，能否实现有效的组合？
• RQ3 一个单一、可扩展的 IR 是否能在不牺牲表达性或可组合性的前提下，实现在多种硬件平台（CPU、GPU、FPGA、加速器）上的高性能？
• RQ4 Relay 在优化和部署具有复杂控制流与异构硬件目标的模型方面，相较于现有框架在多大程度上表现更优？

主要发现

• Relay 在视觉和自然语言处理工作负载上实现了具有竞争力的性能，与最先进框架相当或更优：在 GRU 上比 MxNet 快 2.3 倍，在 TreeLSTM 上快 2 倍。
• 由于能够将复杂控制流编译为单个高度优化的二进制文件，Relay 在 CharRNN 上比 PyTorch 快 1.4 倍，在 TreeLSTM 上快 2 倍。
• 在 ARM 平台（Raspberry Pi 3 和 Firefly RK3399）上的量化推理表现出极小的精度损失（例如，ResNet-18 的准确率从 70.7% 降至 69.4%），同时在激进量化（8/16 位）下获得了显著的性能提升。
• 在基于 FPGA 的 DNN 加速器上，Relay 有效支持了不同硬件设计的评估，显示在内存受限的 NLP 工作负载（如 TreeLSTM）中，批处理可提升吞吐量，但在计算受限的 ResNets 上则无此效果。
• 算子融合与布局转换在 CPU 和 GPU 上均带来显著加速：在 MobileNet 和 ResNet 上最高达 2.3 倍，主要得益于逐元素操作的融合。
• Relay 成功支持多个硬件后端（包括 TPU、Inferentia 和定制 FPGA），展示了强大的可移植性，同时未牺牲优化能力。