0
| 本文作者： 奕欣 | 2017-08-24 05:16 |
数天前，陈天奇团队宣布推出 TVM，在微博上表示，「我们今天发布了 TVM，和 NNVM 一起组成深度学习到各种硬件的完整优化工具链，支持手机，cuda, opencl, metal, javascript 以及其它各种后端。欢迎对于深度学习，编译原理，高性能计算，硬件加速有兴趣的同学一起加入 dmlc 推动领导开源项目社区 。」
据雷锋网AI科技评论了解，大多数现有系统针对窄范围的服务器级 GPU 进行优化，且需要在包括手机、IOT 设备及专用加速器上部署大量工作。而 TVM 是一种将深度学习工作负载部署到硬件的端到端 IR（中间表示）堆栈。也就是说，这类解决方案能够把深度学习模型分发到各种硬件设备上、实现端到端的调优。
它存在三大特点：
能够优化 CPU、GPU 和其他专业化硬件在常规深度学习任务上的计算量；
能够自动转换计算图，使得内存利用率最小化，优化数据布局，将计算模式进行融合。
提供从现有前端到裸机硬件的端到端编译，到浏览器可执行 Javascripts。
雷锋网AI科技评论了解到，TVM 的首篇博客是这样介绍的：
「在 TVM 的帮助之下，开发者只需要少量的额外工作，便可轻易在手机端、嵌入式设备甚至浏览器上运行深度学习任务。TVM 还为多硬件平台上的深度学习工作负载提供统一优化框架，包括依赖全新计算原语的专用加速器。」
而在今天，陈天奇在微博上发布了新的动态，以图森未来胡玉炜的教程介绍着重推介 TVM 的深度学习 op 优化。
「深度学习 op 优化是非常重要但是困难的问题。来自图森未来的胡玉炜写了一个教程介绍了如何利用 TVM 来优化深度学习的 gpu op，通过几十行 python 代码获得比已有 tf 实现两三倍的提升。」
这一文章目前也同步更新到 TVM 博客上，雷锋网 AI 科技评论第一时间做了覆盖和报道。
胡玉炜，北京航空航天大学电子工程硕士，目前 Gap 一年，现在图森未来 HPC 小组实习。这篇文章题为《Optimize Deep Learning GPU Operators with TVM: A Depthwise Convolution Example》（以 Depthwise Convolution 为例，采用 TVM 优化深度学习 GPU 运算符）
高效的深度学习运算符是深度学习系统的核心。通常这些运算符很难优化，需要 HPC 专家付出非常多的努力。TVM 作为一种端到端的张量 IR / DSL 堆栈，能够让整个过程变得更加简单。
这篇文章提供了一个很好的参考，教会开发者如何在 TVM 的帮助下编写高性能 GPU 运算符内核。团队采用的是 Depthwise Convolution（即 topi.nn.depthwise_conv2d_nchw）作为示例，并演示了如何可以改进已经手动优化的 TensorFlow 中的 CUDA 内核。
根据文章的描述，采用 TVM 的最终版本比不同工作负载下的 tf-1.2 中的优化内核快了 2-4 倍，运算符融合速度提升了 3 倍-7 倍。以下是在 GTX1080 下的测试结果， filter size= [1，256,3,3]，stride = [1，1]，padding ='SAME'：
Depthwise Convolution 是一种构建模型的基本思想，能够有效降低深度神经网络的计算复杂度，包括谷歌的 Xception 和 MobileNet 都属于 Depthwise Convolution。
在 TVM 环境下，运行 Depthwise Convolution 的代码如下：
# padding stagePaddedInput = tvm.compute(
(batch, in_channel, height_after_pad, width_after_pad),
lambda b, c, i, j: tvm.select(
tvm.all(i >= pad_top, i - pad_top < in_height, j >= pad_left, j - pad_left < in_width),
Input[b, c, i - pad_top, j - pad_left], tvm.const(0.0)),
name="PaddedInput")# depthconv stagedi = tvm.reduce_axis((0, filter_height), name='di')dj = tvm.reduce_axis((0, filter_width), name='dj')Output = tvm.compute(
(batch, out_channel, out_height, out_width),
lambda b, c, i, j: tvm.sum(
PaddedInput[b, c/channel_multiplier, i*stride_h + di, j*stride_w + dj] * Filter[c/channel_multiplier, c%channel_multiplier, di, dj],
axis=[di, dj]),
name='DepthwiseConv2d')
胡玉炜在文章中提到了优化 CUDA 代码时通常需要注意的三大问题，即数据重用（data reuse）、共享内存（shared memory）和访问冲突（bank conflicts）。
在现代计算架构中，从从内存中加载数据的成本远高于单个浮点计算。因此，我们希望在输入数据被加载到寄存器或 cache 后能够再次使用。
在 depthwise convolution 中有两种形式的数据重用：过滤器重用和输入重用，前者发生在输入通道上滑过并计算多次时，后者在平铺时发生，以 3x3 的 depthwise convolution 为例：
不平铺情况下，每个线程计算 1 个输出元素并加载 3x3 输入数据。16 线程共有 9x16 个负载。
平铺情况下，每个线程分别计算 2x2 个输出元素并加载 4x4 输入数据。4 线程共有 16x4 个负载。
共享内存可以看作 GPU 中的缓存，且是片上的，速度较快。通常的做法是，将数据从全局内存加载到共享内存中，然后块中的所有线程都从共享内存中读取数据。
而为了避免访问冲突，连续的线程最好访问连续的内存地址，如下所示（每种颜色代表一个共享内存库）：
详细内容可参考https://devblogs.nvidia.com/parallelforall/using-shared-memory-cuda-cc/
具体的优化过程
计算填充输入内嵌以节省内存分配
Padding 被明确声明为一个单独的阶段。通过计算内联以避免冗余内存分配：
s = tvm.create_schedule(Output.op)s[PaddedInput].compute_inline()
将一个大通道分成较小的块
一个简单的做法是一个 CUDA 块处理一个输入通道和相应的过滤器，加载到共享存储器后计算：
IS = s.cache_read(PaddedInput, "shared", [DepthwiseConv2d])
FS = s.cache_read(Filter, "shared", [DepthwiseConv2d])
block_y = tvm.thread_axis("blockIdx.y")
block_x = tvm.thread_axis("blockIdx.x")
# bind the dimension of batch (N in NCHW) with block_y
s[Output].bind(Output.op.axis[0], block_y)
# bind the dimension of channel (C in NCHW) with block_x
s[Output].bind(Output.op.axis[1], block_x)
下图为测试结果，在 GTX 1080 上 1000 次运行的平均时间成本，并与 depthwise_conv2d in tensorflow进行比较。
如果是 21 x 21 或者 32 x 32 的通道大小，则性能表现良好，但如果是 64 x 64，那么性能会大大下降。如果做一些修改，那么效果会提升很多：
在一个 cuda 块中实现 32 x 32 的线程，如下：
num_thread_y 和 num_thread_x 这两个参数如何调才能得到最优解？在 Filter = [256, 1, 3, 3] and stride = [1, 1]下：
通过测试，团队得到以下结果：
大规模平铺对于数据重用是有好处的，但不利于本地存储器读取。
num_thread_y 和 num_thread_x 对访问冲突的影响不同。
num_thread_y 和 num_thread_x 的最佳组合，需要实现高效共享内存访问（避免存储区冲突），数据重用和本地内存读取的平衡。
通过强力搜索，在 TVM 中我们可以将 num_thread_y 和 num_thread_x 作为参数传递给 schedule 函数，并尝试所有可能的组合来找到最优组合。
在 TVM 中，Vthread 能够有效支持 Strided Patterns。
在 Filter = [256, 1, 3, 3], stride = [1, 1], blocking_h = 32, blocking_w = 32 的情况下，结果如下：
case 2 比 case 1 更快，因为在 case 2 num_thread_x = 8 和 num_vthread_x = 4 的情况下，确保了连续的线程访问连续的存储器地址，从而避免访问冲突，如下所示（每个颜色表示一个线程的工作负载）：
再来回顾一下和 Tensorflow 的对比：
运算符融合是一个典型优化深度学习网络的方法，在 TVM 中，考虑到原有的模式 depthwise_conv2d + scale_shift + relu，可以稍作如下修改：
生成 IR 如下：
/* Input = [1, 1, 32, 32], Filter = [1, 1, 3, 3], stride = [1, 1], padding = 'SAME' */produce Relu {
// attr [iter_var(blockIdx.y, , blockIdx.y)] thread_extent = 1 // attr [DepthwiseConv2d] storage_scope = "local" allocate DepthwiseConv2d[float32 * 1 * 1 * 4 * 4]
// attr [iter_var(blockIdx.x, , blockIdx.x)] thread_extent = 1 // attr [iter_var(threadIdx.y, Range(min=0, extent=8), threadIdx.y)] thread_extent = 8 // attr [iter_var(threadIdx.x, Range(min=0, extent=8), threadIdx.x)] thread_extent = 8 produce DepthwiseConv2d {
for (i, 0, 4) {
for (j, 0, 4) {
DepthwiseConv2d[((i*4) + j)] = 0.000000f
for (di, 0, 3) {
for (dj, 0, 3) {
DepthwiseConv2d[((i*4) + j)] = (DepthwiseConv2d[((i*4) + j)] + (tvm_if_then_else(((((((1 - di) - i) <= (((blockIdx.x*8) + threadIdx.y)*4)) && ((((blockIdx.x*8) + threadIdx.y)*4) < ((33 - di) - i))) && (((1 - dj) - j) <= (threadIdx.x*4))) && ((threadIdx.x*4) < ((33 - dj) - j))), Input[(((((((((((blockIdx.y + blockIdx.x)*8) + threadIdx.y)*32) + threadIdx.x)*4) + (i*32)) + j) + (di*32)) + dj) + -33)], 0.000000f)*Filter[((di*3) + dj)]))
}
}
}
}
}
for (i2.inner.inner.inner, 0, 4) {
for (i3.inner.inner.inner, 0, 4) {
Relu[((((((((blockIdx.y + blockIdx.x)*8) + threadIdx.y)*32) + threadIdx.x)*4) + (i2.inner.inner.inner*32)) + i3.inner.inner.inner)] = max(((DepthwiseConv2d[((i2.inner.inner.inner*4) + i3.inner.inner.inner)]*Scale[0]) + Shift[0]), 0.000000f)
}
}}
可以看到，每个线程在将 depthwise_conv2d 的结果写入全局内存之前，会计算 scale_shift 和 relu。融合运算符与单个 depthwise_conv2d 一样快。以下是 Input = [1，256，96，96]，Filter = [256，1,3,3]，stride = [1，1]，padding =‘SAME'的结果：
tf-1.2 depthwise_conv2d：251.6 us
tf-1.2 depthwise_conv2d + scale_shift + relu（separate）：419.9 us
TVM depthwise_conv2d：90.9 us
TVM depthwise_conv2d + scale_shift + relu（fusion）：91.5 us
更多优化代码可参考以下链接：
Declare: https://github.com/dmlc/tvm/blob/master/topi/python/topi/nn/convolution.py
Schedule: https://github.com/dmlc/tvm/blob/master/topi/python/topi/cuda/depthwise_conv2d.py
Test: https://github.com/dmlc/tvm/blob/master/topi/recipe/conv/depthwise_conv2d_test.py
雷峰网版权文章，未经授权禁止转载。详情见转载须知。
