NNVM Compiler：AI 框架开放编译器


华盛顿大学 Paul G. Allen 计算机科学与工程学院

亚马逊云服务 AI 团队

DMLC 开源社区

深度学习已经变得无处不在且不可或缺。我们看到在移动电话、GPU、物联网设备和专用加速器等多种平台上部署深度学习工作负载的需求日益增长。上个月，我们发布了 TVM 堆栈，以弥合深度学习框架与面向性能或效率的硬件后端之间的差距。TVM 堆栈使得为深度学习框架构建端到端编译变得容易。但是，我们认为拥有一个适用于所有框架的统一解决方案会更好。

今天，华盛顿大学 Allen 学院和 AWS AI 团队，以及其他贡献者，激动地宣布发布 NNVM 编译器，这是一个开放的深度学习编译器，用于将前端框架工作负载直接编译到硬件后端。我们使用 TVM 堆栈中的两级中间表示 (IR) 构建它。读者可以参考原始 TVM 发布公告，了解有关 TVM 堆栈的更多技术细节。借助 TVM 堆栈，NNVM 编译器可以

  • 在高层图 IR 中表示和优化常见的深度学习工作负载
  • 转换计算图以最大限度地减少内存利用率，优化数据布局并融合不同硬件后端的计算模式。
  • 呈现从前端深度学习框架到裸机硬件的端到端编译管道。

image

NNVM 编译器可以直接从深度学习框架（如 Apache MXNet）获取模型。它还支持模型交换格式，如 ONNX 和 CoreML。ONNX 支持使 NNVM 能够编译来自 PyTorch、Caffe2 和 CNTK 的深度学习模型。CoreML 前端支持将 CoreML 模型部署到非 iOS 设备。

image

优化和部署的分离

image

NNVM 编译器应用图级和张量级优化，并联合优化它们以获得最佳性能。我们采用了与现有深度学习框架不同的方法，现有框架将图优化与部署运行时打包在一起。NNVM 编译器采用来自编译器的传统智慧，将优化与实际部署运行时分离。这种方法提供了大量的优化，但仍然保持运行时的轻量级。编译后的模块仅依赖于最小的 TVM 运行时，在 Raspberry Pi 或移动设备上部署时仅占用约 300KB。

性能

NNVM 编译器仍在积极开发中，我们可以期待未来会有更多改进，但我们已经开始看到有希望的结果。我们对其性能进行了基准测试，并将其与 Apache MXNet 在两种典型的硬件配置上进行了比较：Raspberry Pi 上的 ARM CPU 和 AWS 上的 Nvidia GPU。尽管这两种芯片之间的架构差异很大，但我们可以使用相同的基础设施，只需要更改每种硬件类型的调度即可。

Nvidia GPU

GPU 基准测试和调度由 Leyuan Wang（AWS/UCDavis）和 Yuwei Hu（TuSimple）贡献。我们将 NNVM 编译器与 Apache MXNet（使用 CUDA8 和 cuDNN7 作为后端）在 Nvidia K80 上进行了比较。这是一个非常强的基线，因为 Apache MXNet 开启了自动调优以从 CuDNN 中选择最佳内核。我们还在 MXNet 中使用了优化的深度卷积核来优化 MobileNet 工作负载。

image

可以看出，NNVM 编译器生成的代码优于 K80 上的 Apache MXNet。这些改进归功于联合图级和内核级优化。值得注意的是，NNVM 编译器自行生成所有优化的 GPU 内核，而无需依赖 CuDNN 等外部库。

Raspberry Pi 3b

Rasberry Pi 编译堆栈由 Ziheng Jiang（AWS/FDU）贡献。我们将 NNVM 编译器与 Apache MXNet（使用 OpenBLAS 和 NNPack）进行了比较。我们探索了设置以充分发挥 MXNet 的性能：我们在 NNPACK 中为 3x3 卷积启用了 Winograd 卷积，启用了多线程并禁用了额外的调度线程（因此所有线程都由 NNPack 使用）。

image

可以看出，NNVM 编译器生成的代码在 ResNet18 上快了两倍。MobileNet 上的差距主要是由于现有 CPU DNN 库中缺少深度卷积。NNVM 编译器受益于直接生成高效的 ARM 代码。

致谢

如果没有 DMLC 社区的早期贡献者，这个项目就不可能实现。我们特别感谢 Yuwei Hu（TuSimple）、Leyuan Wang（AWS/UCDavis）、Joshua Z. Zhang（AWS）和 Xingjian Shi（HKUST）对该项目的早期贡献。我们还要感谢所有 TVM 堆栈的贡献者。

在构建 NNVM 编译器时，我们也从以下项目中学习了很多。

  • Theano：可能是最早的深度学习编译器
  • Halide：TVM 使用 HalideIR 作为算术简化和低级降低的数据结构。HalideIR 来源于 Halide。我们在 TVM 中实现降低管道时也向 Halide 学习。
  • Loopy：整数集分析及其循环转换原语的使用。