大V博文系列：PLDI 2021论文分析(三)：DeepCuts-针对GPU的深度学习优化框架

作者：金雪锋

作者主页：https://www.zhihu.com/people/jin-xue-feng

文章来源：https://zhuanlan.zhihu.com/p/393460585

小伙伴们最近分析PLDI一篇很有意思的论文《DeepCuts: A Deep Learning Optimization Framework for Versatile GPU Workloads》，给大家分享一下。

PLDI页面：https://pldi21.sigplan.org/details/pldi-2021-papers/13/DeepCuts-A-Deep-Learning-Optimization-Framework-for-Versatile-GPU-Workloads

论文作者：Wookeun Jung, Thanh Tuan Dao, Jaejin Lee

论文下载链接：https://dl.acm.org/doi/10.1145/3453483.3454038

背景

GPU是目前运行深度学习相关应用的标准后端硬件。目前绝大多数深度学习框架都是使用cuDNN原生高性能库对在GPU上进行的计算做加速优化。然而随着深度学习领域的快速发展，基于cuDNN的框架在GPU优化上并不总能给出最优解。论文给出两个原因：

1. 多变的workloads。cuDNN对于非卷积类神经网络支持不足且无法完美支持所有batch size的输入数据。
2. cuDNN有限的算子融合能力。其只支持少数几种pattern如多个卷积的融合、卷积与biasAdd和ReLU的后向融合，不足以覆盖越来越多样的DL patterns

针对这些限制，一系列深度学习框架（TVM，TensorFlow XLA，TensorRT等）进行了优化并且在部分场景下性能媲美甚至超越cuDNN，但它们仍然无法在各类DL workloads下保持全面稳定的加速效果。其主要原因是获取GPU kernel优化所需信息的难度较大。为提升某个算子在GPU上的运行速度，必须找到使其性能最优的参数组合，包括使用线程块的多少、输入数据的切分大小等。但碍于庞大的搜索空间，想基于所有参数组合去生成代码、评估性能几乎是不可能的。

DeepCuts概述

因此论文作者提出了DeepCuts，一个专为解决深度学习应用中越来越多样的workloads而设计的深度学习框架。其不同于上述其他框架的两个重要创新点：

• 直接使用GPU硬件参数
• 不实际评估每个参数组合的性能，而是通过估计性能上界去筛除理论计算结果较差的参数组合，缩小范围。

对于一个给定的深度学习模型中的计算操作，DeepCuts会搜索产生最优性能GPU kernel的参数组合，并在搜索过程中尝试算子融合，并使用具有后端架构感知的性能估计模型对“definitely-low”的参数组合进行剪枝。基于选中的参数组合，生成数据流图并进一步生成GPU kernel。最后执行这些kernel并为每个算子或融合算子选出使其性能最优的。

主要贡献有：

• DeepCuts在多样的深度学习workload上提供了比cuDNN更具竞争力的性能，且能覆盖小批量推理、大批量推理和训练的场景。作者认为DeepCuts是第一个同时实现高性能和多功能性的框架。
• DeepCuts提出了一种由性能评估模型驱动的代码生成算法，该算法同时考虑了性能优化中的两个关键因素：算子融合和参数搜索。该算法能在相对较短的时间内搜索到使性能最佳的算子融合方式及实现参数。
• 与cuDNN（封闭代码库）不同，DeepCuts在保证性能的同时，算法更加简单易且不依赖任何人工调优。
• DeepCuts评估了其在常用DNN模型包括CNN、RNN、Transformer-based MLP的不同batch size下的推理和训练，及常用硬件后端NVIDIA V100和RTX 2080 GPU上的表现。对比CuDNN，其在小批量推理中分别实现了1.70和1.89倍加速，在大批量推理中分别实现了2.06和2.52倍的加速，在训练中分别实现了1.09和1.10的加速。
• DeepCuts还与TVM、TensorFlow XLA和TensorRT三种现有优化框架进行了比较。其代码生成时间比TVM缩短了13.53×，并且对比三种框架分别实现了1.25、1.48和1.36倍的加速。

相关工作

DeepCuts将主流框架分为两类，并分析了其优缺点。Hand-tuned kernels relied的TensorFlow XLA和TensorRT只在特定算子和规定pattern下能得到较优性能，ML based的TVM (autoTVM) 和 Tensor Comprehensions代码生成策略灵活但是性能不及手写优化。

Table 1对比了DeepCuts与目前较为主流的三大框架并借此表明其在所支持的模型类型、所覆盖的workload类型、涵盖的卷积算法以及性能等多方面的优势。

DeepCuts总体结构

与PyTorch和TensorFlow相似，其输入为描述计算与数据流的无环图，图中的点（node）表示张量计算操作，边（edge）表示数据流向。输出为一系列GPU kernel代码。

四个组成部分：候选者生成器（candidate generator）, 性能评估器（performance estimator）, 代码生成器（code generator）以及代码选择器（code selector）。

首先，候选者生成器将给定workload中的一系列代表张量操作的node拆分为多个非空子集。为每个子集生成单独的GPU kernel（一个子集中的所有单个kernel也会被融合成一个大的kernel）。生成器会生成多种拆分，进行评估并确定性能最优的拆分策略。剩下的三个组成部分，性能评估器（performance estimator）, 代码生成器（code generator）以及代码选择器（code selector）所组成的Evaluate（P）过程既是为了完成每种拆分方式中的每个子集的性能评估。

对于P中的每个子集S，性能评估器会使用性能评估模型找到性能位于top T%的执行参数组合。对于搜索到的每组参数，代码生成器生成相应GPU kernel代码。代码选择器会执行代码，记录运行时间，选出性能最佳的kernel并更新P的总运行时间。

分割算法：基于当前拆分仍有加速潜力的前提下，DeepCuts增量地生成更多拆分，直至没有新的拆分可生成。尽管无法枚举所有可能的候选拆分，DeepCuts的作者认为其拆分算法可以枚举绝大多数对性能有提升的拆分方案。举例来说，对于X, Y, Z三个连续子集，只有当融合其中的两个子集能够获得性能收益时，此算法才会将三者融合成一个。因为当X+Y和Y+Z都会导致性能劣化时，X+Y+Z的融合几乎不会获得性能提升。

Table 2列出了DeepCuts所能覆盖处理的节点类型，分为复杂操作（计算密集且耗时的操作如卷积和矩阵乘法）和简单操作（计算复杂度和输入数据大小成正比的操作）。后面会介绍DeepCuts基于不同算子类型给出的融合策略。

性能评估模型

1、kernel的实现参数

Table 3和Table 4列出了DeepCuts的性能评估模型中涉及到的实现参数。Table 3以卷积操作为例，给出了由输入tensor定义的固定参数。Table 4则给出了变量参数，DeepCuts在优化过程中需要找到的能够给出较优性能的参数组合中的参数就是这些变量参数。DeepCuts将生成的代码分为取数据和计算两个阶段。

为缩小搜索空间，DeepCuts将这些实现参数限制为：每个线程计算的各维度上数据为2的指数幂，每个线程块处理的则是每个线程处理的对应维度数据的整数倍。

2、性能限制因素

DeepCuts给出四个主要限制kernel性能且与GPU架构相关的因素，并设计了四个相对应的metric值（ , , , ）。通过计算这些值的上界，对理论上不可能有较优性能的kernel进行剪枝。相关参数列在Table 5到8中。

下面对这四个因素逐个阐述：

• 全局内存带宽global memory bandwidth（涉及变量见Table 6, 8）

一个线程块的计算密集程度（the operational intensity of a thread block）：

并计算出该因素相应的度量值：当前OI和理论峰值的比值，

若GMRatio=1，表明全局内存带宽并不是限制性能的瓶颈。

• 共享内存延迟shared memory latency（涉及变量见Table 6, 7）

DeepCuts根据一个线程内计算操作次数与共享内存加载操作次数的比值去估算共享内存的延迟。

其中COEFbc代表共享内存中产生bank conflicts的程度，可根据共享内存中每个线程束（warp）访问的数据地址计算得到。举例，若一个线程在每次共享内存加载数据的时间内能进行10次浮点数运算，且已知共享内存延迟为20个cycles并且不存在bank conflicts，那么可以得到每20个cycles中有10个cycle可被隐藏，则SMRatio=0.5；若能进行20次浮点数运算，则SMRatio=1，表明共享内存延迟也不是限制性能的因素。

• 流处理器间工作负荷的不平衡workload imbalance between streaming multiprocessors (SMs) （涉及变量见Table 4, 5）

若每个线程块所处理的数据个数相同，则线程块间不存在工作负载不均。但当每个线程块所处理的数据不是SM个数的整数倍时，这种不平衡就会显现。使用的线程块数量以及相应的度量值计算如下：

• 硬件资源限制limitation of hardware resources（涉及变量见Table 4, 5）

这项指标是为了筛除那些所需资源超出现有可用硬件资源导致生成代码无法运行的kernel。不同于以上三个因素的度量值都是处于0到1之间的变化值，COEFr只有0和1两种取值。

3、估算性能上界

根据上述四个度量值，DeepCuts给出一个总的度量值PUL，来估算一个张量操作的性能上界：

PUL表示对于目标GPU，当前张量操作的最高性能与理论峰值的比值。DeepCuts会计算每个可能出现的参数组合的PUL，只选取结果在前1%的组合去进行下一步的代码生成。对较为常用的张量操作进行了大量参数组合的实验后（ResNet中的3*3卷积及BERT中的矩阵乘法，实验大于十万组参数），测试结果证明性能评估模型没有错误地筛除掉任何性能最优的参数组合。

4、共享内存层级与寄存器层级的融合

寄存器层级融合：两个简单算子融合成一个简单算子（element-wise OP + ReLU），或一个简单算子和一个复杂算子融合成一个复杂算子（convolution + bias add）。其PUL的计算遵循如下规则：

共享内存层级融合：当后向算子的多个线程都需要使用前向算子中一个线程的计算结果（convolution + convolution）。其PUL计算限制如下：

数据流图生成

生成数据流图（data-flow- graph, DFG）包含以下三个部分：

• 基本DFG生成 baseline DFG generation：

表示一个线程块中的计算。当给出这样一组参数S，

基本DFG生成如下图Figure 3：

当C轴需要被切分（即C input < C）时，代码生成器会生成包含取数据与计算交替的循环，称为looped complex operation。此时DeepCuts会生成两个不同的DFG，一个表示循环体内的操作，另一个表示退出循环时的操作（如将计算结果存储到共享内存中）。

• DFG拼接 DFG concatenation

Figure 4展示了基于共享内存层级和寄存器层级的两种融合方式下产生的不同DFG。

• 每个线程的子图提取subgraph extraction for a GPU thread

DeepCuts将DFG中最后一个共享内存的存储操作切分为多个等量的数据块，并分配给不同的线程，然后顺着图中的边（edge）回溯地标记所有访问到的点。到达没有前向节点的节点后，山区所有未被标记的节点即可获得每个线程的sub-DFG。

GPU-kernel代码生成

DeepCuts基于每个线程的sub-DFG生成kernel代码。首先发射kernel函数头文件和线程/线程块ID的获取，然后根据性能评估器中选定的参数生成从全局内存到共享内存的取数据代码。最后对于计算部分，根据反向逆序遍历sub-DFG然后为每个节点发射代码。

对于looped complex operation, 代码生成器也会生成循环结构，以及循环体内的取数据与计算两块代码。常见的优化手段，双缓冲区或称为数据预取，也是在这一步实现，生成器会生成normal和prefetch两块代码，prefetch就是用于隐藏全局内存的加载延迟。

DeepCuts还进行了如下对于共享内存的优化：

基于内存的索引函数Memory-based index function:由于CUDA代码的SPMD形式，为使得每个线程以数据并行的形式正确访问到其相应数据元素，用thread ID表示该线程在数据中访问的索引。索引函数使用两种计算方式：Computation-based方式是在线程访问相应数据元素时计算索引，memory-base则是在host侧提前计算好并储存在共享内存中。

向量化I/O指令：Exploiting vector I/O instructions:即同时从共享内存中加载多位数据如常见的128-bit load，可通过减少加载的指令个数而提高性能。

数据补齐减少存储体冲突Bank-conflict aware padding: 通过对输入数据的补零使数据分布对齐从而减少共享内存中的存储体冲突。

评估

GPU架构：NVIDIA Volta (V100) 和Turing (RTX 2080)

Benchmark模型：

评估方式：包含每个benchmark模型中训练和推理的workload。对于非CNN模型：记录1个epoch的执行时间，评估不同batch size下1个训练和4个推理共20种workload；对于CNN模型：记录100个mini-batch的执行时间，评估不同batch size下1个训练和2个推理共12种workload。排除了文件I/O及CPU-GPU间数据传输的影响。

评估框架：DeepCuts-no-fusion, DeepCuts-fusion, DeepCuts-mixed（cuDNN/cuBLAS与DeepCuts生成的kernel同时使用）。

对比框架：TVM-default, TVM-autotvm; TensorFlow, TensorFlow-XLA, TensorRT-integrated TensorFlow; PyTorch.

从Table 11和12中看出，在32个workload中的大部分DeepCuts-mixed都能获得最好性能，平均性能优于其余框架。DeepCuts-fused在不适用cuDNN原生库的情况下，也能在其中9个workload中打败其他框架。论文中还强调，DeepCuts-mixed中只有25%的算子使用的是cuDNN/cuBLAS原生库，其余75%仍是框架自主生成的代码。

根据Figure 5可分析得出，DeepCuts在不同的benchmark模型场景下的表现也略有不同，但总体都能达到比其他框架更优的性能。

Figure 6拆分了不同种类算子的执行时间并进行了不同框架下的对比。得益于DeepCuts在卷积和矩阵乘法等复杂算子上较优的代码生成，以及在简单算子上的融合优化，其总执行时间是相对较短的。

Figure 7则给出了代码生成的时间。TVM-autotvm的代码生成时间很慢且波动较大。TensorFlow-XLA和TensorRT使用已有的cuDNN/cuBLAS原生库或事先写好的kernel因此代码生成几乎不费时间，生成速度是DeepCuts的1058倍和652倍。DeepCuts的生成速度是TVM-autotvm的13.53倍，不过其代码生成时间与batch增大引起的搜索空间增大成正比增长。

未来工作

支持算子多样性：DeepCuts尝试为Shift-Conv-Shift-Conv和Octave Convolution（需拆分为两个kernel）这两种不常见的算子进行kernel代码生成。未来希望通过实现水平融合进一步提高性能。

支持GPU后端多样性：目前DeepCuts只支持NVIDIA GPU作为后端。未来希望拓展给予Open-CL的DeepCuts至支持AMD GPUs以及更多移动端侧设备。

结论

本文介绍了DeepCuts，一个用于多样化GPU工作负载的深度学习优化框架。其主体流程如下：对于给定DL workload中的一系列计算操作，DeepCuts在考虑算子融合的同时，搜索对于每个算子能给出最佳性能的实现参数组合，并使用性能估计模型对会导致“绝对慢”的case进行剪枝。然后，它使用选定的参数组合生成数据流图并生成GPU kernel，通过执行kernel并对比执行时间，为当前workload中的每个算子或融合算子确定使其性能最佳的执行参数。

实验评估表明，对于DL工作负载，对于CNN、RNN和Transformer-based MLP的训练和推理，DeepCuts的性能与cuDNN相当甚至有所超越。当DeepCuts在kernel选择过程中同时包括cuDNN原语和自生成代码时（即DeepCuts-mixed）：在NVIDIA V100 GPU上，它比cuDNN和PyTorch有1.13和1.46倍的加速；它在推理上与TVM相比实现1.24倍加速（TVM不支持训练），且代码生成时间减少了13.53×；它与支持训练的TensorFlow-XLA相比实现了1.38倍的加速；DeepCuts在与NVIDIA V100体系结构不同的NVIDIA RTX 2080 GPU上也实现了良好的性能。作者还表示DeepCuts的代码也将开源。

点评

DeepCuts作为新生的深度学习优化框架，其大体流程如数据流图生成、算子融合、融合后的代码生成及生成过程中的代码优化等环节，与其他框架或算子库，包括TVM,TensorFlow, TensorRT以及MindSpore的图算融合&AKG，是类似的。

其值得借鉴的两大思路，一个是将我们常见的tuning思路扩展至整个框架上，包括算子融合方式和执行参数选择（目前tuning多用于仅仅选择参数），并且是结合了先剪枝、后逐个执行的策略，以在相对短的时间内获得较优性能。

另一个是DeepCuts针对GPU后端会对性能产生限制的四个因素设计了简洁明确的metric value及相应计算公式，以计算出每套执行参数的性能上界并帮助选择。其对GPU后端参数的直接使用（可应对硬件型号调整），清晰的评估逻辑以及对各因素影响的系统性量化，这些特点是值得借鉴的。

从另一方面来说，DeepCuts也具有一定局限性。DeepCuts选择最终融合策略与执行参数的方式仍然是通过逐个遍历执行（尽管进行过剪枝），因此其在大batch场景下算子生成时间仍然较长。

另外论文中给出的各部分算法较为简单，覆盖场景较少。首先其partition算法只考虑将多个单算子融合成一个大算子（主要适用于简单算子），而没有考虑将单个复杂算子拆分成多个算子再重新融合的场景；其次，文中给出的性能评估metrics所能应对的场景也有所局限，比如对于融合算子，其只能处理将相邻算子纵向拼接的简单融合，同时还对前后向算子shape有一定要求；同时，整篇文章只提及单精度计算操作，而没有对半精度的计算场景做优化，也就意味它没有使能NVIDIA GPU专为半精度和混合精度计算设计的Tensor Cores加速单元。可能DeepCuts目前还难以处理应用Tensor Cores带来的更大的参数搜索空间，和更复杂的代码生成流程。

总的来说，DeepCuts的框架实现思路能够为深度学习框架后续发展带来不少启发，不过有些部分的算法还有待丰富，可以期待一下代码的开源。

注：文章中所有图片均来自论文《DeepCuts: A Deep Learning Optimization Framework for Versatile GPU Workloads》。