[论文解读] High-Performance Deep Learning via a Single Building Block

Evangelos Georganas, Kunal Banerjee+7 more|arXiv (Cornell University)|Jun 15, 2019

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

一句话总结

本文提出批量归约GEMM核作为高性能深度学习工作负载的通用构建模块，仅用3,000行高级代码即可高效实现RNN、CNN和MLP。通过优化此单一核，作者在CPU上实现的性能超过厂商优化库，在GPU上达到或超越自动调优的核，展示了可扩展、可维护的专用低层核的替代方案。

ABSTRACT

Deep learning (DL) is one of the most prominent branches of machine learning. Due to the immense computational cost of DL workloads, industry and academia have developed DL libraries with highly-specialized kernels for each workload/architecture, leading to numerous, complex code-bases that strive for performance, yet they are hard to maintain and do not generalize. In this work, we introduce the batch-reduce GEMM kernel and show how the most popular DL algorithms can be formulated with this kernel as the basic building-block. Consequently, the DL library-development degenerates to mere (potentially automatic) tuning of loops around this sole optimized kernel. By exploiting our new kernel we implement Recurrent Neural Networks, Convolution Neural Networks and Multilayer Perceptron training and inference primitives in just 3K lines of high-level code. Our primitives outperform vendor-optimized libraries on multi-node CPU clusters, and we also provide proof-of-concept CNN kernels targeting GPUs. Finally, we demonstrate that the batch-reduce GEMM kernel within a tensor compiler yields high-performance CNN primitives, further amplifying the viability of our approach.

研究动机与目标

• 解决在多种架构上为多样化深度学习工作负载实现低层核优化时面临的组合爆炸问题。
• 通过将所有主要深度学习原语统一为单一高度优化的核，降低深度学习库开发的复杂性。
• 仅通过围绕单一优化核进行循环调优，实现在RNN、CNN和MLP训练与推理中均达到高性能。
• 证明单一高度优化的核可超越CPU和GPU上的手工调优、厂商优化实现。
• 通过与TVM等张量编译器集成，实现高性能可移植的深度学习原语。

提出的方法

• 提出批量归约GEMM核，该核将输入子张量块的批量乘法运算结果归约到单个输出子张量块。
• 将所有主要深度学习原语（RNN/LSTM、CNN、MLP）表述为该核的组合，从而抽象掉底层复杂性。
• 针对CPU和GPU架构对批量归约GEMM核进行优化，采用细粒度内存访问、预取和数据布局变换。
• 仅通过围绕优化核进行循环调优，在3,000行高级代码中实现深度学习原语。
• 将该核集成至分布式训练框架（如GNMT和ResNet-50），并在多节点集群中展示其可扩展性。
• 在TVM张量编译器中使用批量归约GEMM核原型实现CNN核，性能达到与自动调优原语相当的水平。

实验结果

研究问题

• RQ1 单一高度优化的核是否足以替代在不同深度学习工作负载和架构上所需的数百个专用手工调优核？
• RQ2 通过将所有原语抽象为单一核，能否显著提升深度学习库在性能可移植性和可维护性方面的表现？
• RQ3 统一核方法是否能在CPU和GPU工作负载中超越厂商优化的、临时实现的代码？
• RQ4 批量归约GEMM核在分布式多节点环境中实现高性能推理与训练方面的有效性如何？
• RQ5 张量编译器能否利用该核实现与自动调优、厂商优化原语性能相当的性能？

主要发现

• 在Xeon Skylake-SP CPU上，批量归约GEMM核达到峰值性能的83%，超过通用GEMM方法（61%和49%），并以2%更高的效率匹配厂商优化的mkl-dnn（81%）。
• 基于批量归约GEMM核的作者实现的CNN原语在CPU上比厂商优化的mkl-dnn库快1.24倍，在多节点集群上快1.4倍。
• 在集成GPU上，该核性能与厂商提供的CNN核相当，证明其架构无关性。
• 在多节点CPU集群上，GNMT和ResNet-50的端到端训练性能超过厂商优化库最高达2.3倍。
• 在基于TVM的原型中，批量归约GEMM核实现了与自动调优、Amazon-AutoTVM优化代码性能相当的CNN原语。
• 整个深度学习库实现仅需3,000行高级代码，相比现有库动辄数万行代码，显著降低了代码库复杂度。