使用自动代码生成技术TVM优化深度学习算子的一些思考 - 知乎



随着人工智能的日益普及，深度学习的新算法和新模型层出不穷，以tensorflow、pytorch、mxnet等为代表的深度学习框架也在持续更新，以期在最大程度上支持新的特性。灵活性和高性能是现代深度学习训练框架的两大追求。灵活性主要是指python语言的灵活性，这种灵活性能够帮助研究员快速实现新的算法或者提出新的神经网络模型。由于python提供了比较大的灵活性，研究员在实现算法时会使用大量细碎算子，而且这种细碎算子很多都是非计算类的，例如转置、zip、slice、irregular split/concat和topk等。细碎算子在模型训练过程占用的时间越来越多。一般而言，分类模型中细碎算子较少，而检测等模型细碎算子占用时间甚至超过50%。因此，训练框架的高性能很大程度上依赖于对这些连续细碎算子的整体优化。
需要指出的是，我们这里说的算子优化，并不只是对于单个op的优化，更多的是一个python中定义的函数，或者说一系列单个操作的集合。对于类似于卷积和矩阵乘法这样的单个基本算子，在诸如cudnn、cublas中已经有很高效的实现，目前的深度学习框架基本都已经使用了这些高效的计算库, 因此单纯的优化这些算子对于训练模型整体的速度提升是很有限的。当然，这些通用计算库因为要适应不同的输入规模，在实际使用中效果不一定总是很好。例如，cudnn的reduction函数在(n,c,h,w)分别规约成(n,c,1,w)和(n,c,h,1)时性能就有很大的差别，研究员可能就需要自己实现更高效的规约函数。
下面以mmdetection里的delta2bbox函数为例，说明算子手动优化和深度学习框架的实现相比，为什么能够获得加速。以pytorch框架为例，基本上python代码的每一行，都会调用pytorch计算库中相应的gpu kernel。 利用pytorch提供的trace工具追踪这个函数，我们发现这些gpu kernel被调用：repeat、mul、add、slice、clamp、unsqueeze、exp、addcmul、 stack。这些连续细碎算子，对于程序编写是有利的，但是GPU连续调用这些细碎算子，会使得其在利用率不高的情况下连续执行较长时间(尤其当算子是非计算型或者计算规模较小时)，整体上表现就是程序运行时间长，gpu占用率较低。而通过将这个函数实现为一个或者两个连续的kernel，就可以充分发挥GPU的计算能力。其实这也是当前类似于TensorRT这样的推理框架进行加速的一个重要手段，在宏观上，对计算图进行表达优化，在微观上，进行算子层面的融合计算，再配合高效和多版本的候选kernel，加上量化以及模型轻量化等手段，就可以获得数十倍乃至上百倍的加速比。因此，从这个角度来讲，一个深度学习框架，能够在满足硬件约束的情况下自动融合更多的op, 并且能在底层有更优秀的计算库支撑，其整体性能就会更优。

  1.     def delta2bbox(rois，deltas, means=[0, 0, 0, 0],
  2.                stds=[1, 1, 1, 1], max_shape=None,
  3.                wh_ratio_clip=16 / 1000):
  4.     means = deltas.new_tensor(means).repeat(1, deltas.size(1) // 4)
  5.     stds = deltas.new_tensor(stds).repeat(1, deltas.size(1) // 4)
  6.     denorm_deltas = deltas * stds + means
  7.     dx = denorm_deltas[:, 0::4]
  8.     dy = denorm_deltas[:, 1::4]
  9.     dw = denorm_deltas[:, 2::4]
  10.     dh = denorm_deltas[:, 3::4]
  11.     max_ratio = np.abs(np.log(wh_ratio_clip))
  12.     dw = dw.clamp(min=-max_ratio, max=max_ratio)
  13.     dh = dh.clamp(min=-max_ratio, max=max_ratio)
  15.     px = ((rois[:, 0] + rois[:, 2]) * 0.5).unsqueeze(1).expand_as(dx)
  16.     py = ((rois[:, 1] + rois[:, 3]) * 0.5).unsqueeze(1).expand_as(dy)
  17.     pw = (rois[:, 2] - rois[:, 0] + 1.0).unsqueeze(1).expand_as(dw)
  18.     ph = (rois[:, 3] - rois[:, 1] + 1.0).unsqueeze(1).expand_as(dh)
  20.     gw = pw * dw.exp()
  21.     gh = ph * dh.exp()
  22.     gx = torch.addcmul(px, 1, pw, dx)
  23.     gy = torch.addcmul(py, 1, ph, dy)
  25.     x1 = gx - gw * 0.5 + 0.5
  26.     y1 = gy - gh * 0.5 + 0.5
  27.     x2 = gx + gw * 0.5 - 0.5
  28.     y2 = gy + gh * 0.5 - 0.5
  29.     if max_shape is not None:
  30.         x1 = x1.clamp(min=0, max=max_shape[1] - 1)
  31.         y1 = y1.clamp(min=0, max=max_shape[0] - 1)
  32.         x2 = x2.clamp(min=0, max=max_shape[1] - 1)
  33.         y2 = y2.clamp(min=0, max=max_shape[0] - 1)
  34.     bboxes = torch.stack([x1, y1, x2, y2], dim=-1).view_as(deltas)
  1.  下面结合个人的经验，以GPU平台为例，谈一些算子优化的体会。1）op融合；2）借助高效的计算库；3）自动代码生成和auto-tuning；4）手动实现。<br />1）op融合。如前所述，个人认为op自动融合是深度学习框架进一步提升性能的关键，是未来一个重要的研究方向，也是运行时代码生成和优化需要重点关注的地方。在实践中，笔者发现除了少部分手动实现的单个op会比pytorch的框架更加快速，多数单个op性能已经足够好。手动优化后超额的加速比在很大程度上是由于op融合产生的。以上面的代码为例，dw和dh在计算时的slice操作和clamp操作就可以合并。op融合的另一个好处是减少内存申请和释放的操作次数，有助于storage flatten, 当然op融合由于在一个kernel中进行大量的操作，某些情况下可能会增大内存或者显存的消耗，但这有助于提高GPU的利用率。op融合除了依赖于训练框架本身或者是性能优化工程师，也依赖于研究员对函数或者模块的设计与实现。在深度学习领域，一个比较好的例子就是conv\batch norm\relu的融合，或者我们简单的称其为CBR法则。<br />2）借助于高效的计算库。计算芯片的成功，在很大程度上依赖于其生态和计算库。对于性能优化工程师而言，除了要了解计算芯片的硬件架构、编程模型和语言，熟练的使用其计算库也是必不可少的一项技能。由于深度学习涉及到的主要是tensor之间的各种操作，因此主要用到的计算库有cuDNN、cuBLAS、Thrust、cuRAND等。cuDNN是用于深度神经网络的GPU加速库。它强调性能、易用性和低内存开销，支持常见的DL操作，而且在持续加入新特性。cuDNN已经被集成到更高级别的机器学习框架中，如Tensorflow、Pytorch等。简单的插入式设计理念可以让开发人员专注于神经网络模型的高效快速实现，而不是持续调整性能。Thrust是一个类似于C++ STL的计算库，提供host和device两个版本的实现，其sort、reduce、transform、copy、permute以及zip等操作高效而易于使用，目前thrust已经被集成到pytorch中，其部分排序类算法的底层就是基于Thrust 的API实现的。当然，通用计算库的性能不总是足够好的，需要根据实际情况灵活选用。<br />3）自动代码生成和auto-tuning。由于深度学习中大量涉及较为规则的张量计算，近年来，深度学习编译器开始出现并受到重视。其中，比较有代表性的一个就是TVM。TVM是一个用于深度学习系统的编译器堆栈。它旨在缩小以生产力为中心的深度学习框架与以性能和效率为中心的硬件后端之间的差距。TVM可以与深度学习框架合作，为不同的后端提供端到端编译。TVM既可以进行离线时的算子优化，也可以进行运行时代码生成和优化。这里，我们只讨论前者，也就是离线算子级别的优化。通过TVM自动代码生成的功能，使用python接口，可以部分提升优化人员的工作效率。对于较为规则的算子，TVM已经可以生成效率很高的代码。不过，单从算子优化这一层级考虑，由于tvm主要针对于嵌套循环的优化，这就使得其应用范围受到很大的局限性，类似topk、nms这样的操作，就很难用tvm进行描述。同时，其封装的部分high level API，由于只是采用默认的schedule，可以说是只有计算描述，没有高效的schedule设计，生成的GPU kernel只能用作参考。用户需要手动添加schedule， 或者在低效率的代码基础上进行手动修改。在实际使用中，个人感觉TVM可以考虑进行的改进操作如下：
    1. 支持inplace操作，这样可以增大对于element-wise这一类没有数据依赖性操作的支持；
    1. 支持形如a[b[i][j]][c[i][j]]的操作，也就是说lamada表达式不是结果数组a的直接索引，a的索引由另外的数组b和c给出，从而适应对于张量部分赋值(slice)这一类常见操作的适应性；
    1. 支持concat\split\reshape这三种只改变张量索引形式和循环变量范围，但是数据在内存或者显存中没有变化的操作；
    1. 支持zip操作，zip的多个tensor可能具有不同的数据类型，使得zip中张量可以进行互操作，或者调用同样的外部算法；
    1. auto-tuning功能支持对于具有相同shape的多个output的自动调优，使得auto-tuning推荐的config对于多输出的kernel更加精确；
    1. 离线和运行时生成代码的可用性。就个人体会而言，TVM目前的离线代码生成过于简化，而运行时代码过于复杂。拿离线代码生成来说，当给定Tensor形状时，生成的代码过于简单，用户直接使用很容易丢失一部分的区间判断，造成代码逻辑错误，而运行时生成的代码可读性又很差，过于庞杂，目前的一个实用小技巧是通过互不相同的质数给定shape后手动进行变量替换。

4）手动实现。关于各种硬件后端的程序优化，已经有很多详尽的资料。个人感觉，程序优化就是选用最合适的算法，充分发挥各种计算芯片分级存储体系的潜力，在加上合适的任务调度和分配方式，从而最大化数据吞吐量，尽可能发挥计算芯片的计算能力。从目前的情况来看，CPU + 其他芯片的异构计算体系，应该是一个趋势。
随着自动代码生成工具功能的日益完善，各种高性能计算库的不断更新，以及各种计算芯片的推出，算子优化会有越来越多的选择和工具链，这对算子优化人员来说，应该是一件好事。