将 TVM 融入 TensorFlow，在 GPU 上实现全面提速

雷锋网 AI 研习社按，日前，阿里机器翻译团队和 PAI 团队发表博文，阐述将 TVM 引入 TensorFlow，可以带来至少 13 倍的 batch 矩阵相乘（matmul）加速。雷锋网 AI 研习社将原文编译整理如下：

背景

神经机器翻译（NMT）是一种端到端的自动翻译方法，可能克服传统的基于短语的翻译系统的缺点。最近，阿里巴巴集团正致力于在全球电子商务中部署 NMT 服务。

目前，我们将 Transformer（ https://arxiv.org/pdf/1706.03762.pdf ） 作为 NMT 系统的核心组成。相较于传统基于 RNN/LSTM 的方法，它更适合于高效的离线训练，有着相同或更高的精度。

Transformer 在时间步长中打破了相关性，对离线训练更友好，但在在线推理上，它并没有那么高效。我们在生产环境中发现初版 Transformer 的推理速度大约比 LSTM 版本慢 1.5 倍到 2 倍。为了提高推理性能，我们已经进行了一些优化，包括图级别的 op 融合、循环不变节点外提（loop invariant node motion）。我们观察到一个特殊问题：batch 矩阵相乘是 Transformer 中的一个关键问题，目前它在 cuBLAS 中的实现并未得到很好的优化。

图1：Transformer 模型架构

下图表明，通过 TVM 生成的内核可以带来至少 13 倍的 batch 矩阵相乘加速，伴随算子融合，速度将更快。

batch 矩阵相乘

为什么选择利用 batch 矩阵相乘

在 Transformer 中，batch 矩阵相乘被广泛应用于 multi-head attention 的计算。利用 batch 矩阵相乘，可以并行运行 attention 层中的 multiple heads，这有助于提高硬件的计算效率。

图2：左图为 Scaled Dot-Product Attention，右图为并行运行若干 attention 层的 Multi-Head Attention

我们在推理阶段对 Transformer 模型进行了全面分析，结果表明，batch 矩阵相乘计算的开销达到 GPU 内核执行时间的 30%。当使用 nvprof 对 cuBLAS batch 矩阵相乘内核做一些第一原理（first-principle）分析，很明显，这种方法的表现并不好，同时我们还发现几个有趣的现象。

什么是 batch 矩阵相乘

通常，batch 矩阵相乘计算会在一批矩阵上执行矩阵-矩阵乘法。batch 被认为是「统一的」，即所有实例都具有相同的维度（M，N，K）、leading 维度 (lda，ldb，ldc) 和它们各自的 A、B、C 矩阵的转置。

batch 矩阵相乘计算具体可以描述如下：

void BatchedGemm(input A, input B, output C, M, N, K, batch_dimension) {
 for (int i = 0; i < batch_dimension; ++i)  {
   DoGemm(A[i],B[i],C[i],M,K,N)
 }
}

batch 矩阵相乘形状

在语言翻译任务中，batch 矩阵相乘的形状比在其他工作负载下的常规矩阵相乘计算要小得多。Transformer 的形状与输入语句的长度和解码器步长有关。一般来说小于 30。

至于 batch 维度，当给定推理 batch 大小时，它是固定数字。例如，如果 batch size 是 16，beam size 是 4，batch 维度是 16 * 4 * #head (在 multi-headattention 中 head 的数目，通常为 8)。矩阵 M、K、N 的范围在 [1, max decode length] 或 [1, max encode length] 内。

batch 矩阵相乘的性能问题

首先，我们在理论上对 batch 矩阵相乘内核进行了 FLOP 分析。结果非常有趣：所有 batch 矩阵相乘的计算强度都是受限的（TFLOP 数少于 1）。

然后，我们通过 nvprof 描述了多形状 batch 矩阵相乘的 cuBLAS 性能。下面的表格中是使用 NVIDIA M40 GPU（CUDA 8.0）得到的一些指标。

即使形状不同（在 M、N、K 间变化），所有 maxwell_sgemmBatched_128x128_raggedMn_tn 调用执行的都是相同的 FLOP 数，这比理论值大得多。从中可以推断，所有这些不同的形状最终都会被填充成确定的形状。在所有的形状中，即使在最好的情况下，理论 FLOP 只占实际执行 FLOP 的 2.74%，因此大多数计算都是多余的。类似地，调用另一个 cuBLAS 内核 maxwell_sgemmBatched_64x64_raggedMn_tn 也出现相同情况。

显而易见，cuBLAS batch 矩阵相乘的执行效率很低。基于这个原因，我们在 NMT 中使用 TVM 生成高效的 batch 矩阵相乘内核。

batch 矩阵相乘计算

在 TVM 中，普通 batch 矩阵相乘计算声明如下：

# computation representation
A = tvm.placeholder((batch, M, K), name='A')
B = tvm.placeholder((batch, K, N), name='B')
k = tvm.reduce_axis((0, K), 'k')
C = tvm.compute((batch, M, N),
        lambda b, y, x: tvm.sum(A[b, y, k] * B[b, k, x], axis = k),
        name = 'C')

调度优化

在声明计算之后，我们需要仔细地设计调度来更好地发挥性能。

调节 block/线程数的参数

# thread indices
 block_y = tvm.thread_axis("blockIdx.y")
 block_x = tvm.thread_axis("blockIdx.x")
 thread_y = tvm.thread_axis((0, num_thread_y), "threadIdx.y")
 thread_x = tvm.thread_axis((0, num_thread_x), "threadIdx.x")
 thread_yz = tvm.thread_axis((0, vthread_y), "vthread", name="vy")
 thread_xz = tvm.thread_axis((0, vthread_x), "vthread", name="vx")

 # block partitioning
 BB, FF, MM, PP = s[C].op.axis
 BBFF = s[C].fuse(BB, FF)
 MMPP = s[C].fuse(MM, PP)
 by, ty_block = s[C].split(BBFF, factor = num_thread_y * vthread_y)
 bx, tx_block = s[C].split(MMPP, factor = num_thread_x * vthread_x)
 s[C].bind(by, block_y)
 s[C].bind(bx, block_x)
 vty, ty = s[C].split(ty_block, nparts = vthread_y)
 vtx, tx = s[C].split(tx_block, nparts = vthread_x)
 s[C].reorder(by, bx, vty, vtx, ty, tx)
 s[C].reorder(by, bx, ty, tx)
 s[C].bind(ty, thread_y)
 s[C].bind(tx, thread_x)
 s[C].bind(vty, thread_yz)
 s[C].bind(vtx, thread_xz)

我们融合了 batch 矩阵相乘的外部维度，例如 op 维的 BB 和 FF 在 batch 矩阵相乘计算中通常称为「batch」维，我们用一个因子 (number_thread * vthread) 分割外部和内部维度。

在 batch 矩阵相乘中不需要 Strided 模式，因此将虚拟线程数（vthready 和 vthreadx）都设置为 1。

找到 number_thread 的最佳组合

下面的结果是基于 NVIDIA M40 GPU（CUDA 8.0）。

基于过去的经验，找到 num_thread_y 和 num_thread_x 最佳组合的方法是通过暴力搜索（brute-force search）。经过暴力搜索后，可以找到当前形状的最佳组合，在当前的计算中，num_thread_y = 8，num_thread_x = 32。

将 batch 矩阵相乘与其他运算融合

现有的「黑盒」cuBLAS 库调用一般会作为常用的「op 融合」优化策略的边界。然而，利用生成的高效 batch 矩阵相乘内核，融合边界极易被打破，将不仅仅是各个元素之间的融合，因此可以获得更好的性能改进。

从计算图中可以看出，batch 矩阵相乘之后总是伴随着广播加法运算或转置运算。

通过将「加法」或「转置」运算与 batch 矩阵相乘融合，可以减少内核启动开销和冗余内存访问时间。

batch 矩阵相乘和广播加法融合计算的声明如下：

# computation representation
A = tvm.placeholder((batch_size, features, M, K), name='A')
# the shape of B is (N, K) other than (K, N) is because B is transposed is this fusion pattern
B = tvm.placeholder((batch_size, features, N, K), name='B')
ENTER = tvm.placeholder((batch_size, 1, M, N), name = 'ENTER')
k = tvm.reduce_axis((0, K), 'k')
C = tvm.compute(
          (batch_size, features, M, N),
          lambda yb, yf, m, x: tvm.sum(A[yb, yf, m, k] * B[yb, yf, x, k], axis = k),
          name = 'C')
D = topi.broadcast_add(C, ENTER)

batch 矩阵相乘和转置融合计算的声明如下：

# computation representation
A = tvm.placeholder((batch_size, features, M, K), name='A')
B = tvm.placeholder((batch_size, features, K, N), name='B')
k = tvm.reduce_axis((0, K), 'k')
C = tvm.compute(
          (batch_size, M, features, N),
          lambda yb, m, yf, x: tvm.sum(A[yb, yf, m, k] * B[yb, yf, k, x], axis = k),
          name = 'C')

融合内核性能

测试生成代码性能时，形状选择为 [batch=64, heads=8, M=1, N=17, K=128]。选择 17 作为序列长度是因为它是我们生产中的平均输入长度。

• tf-r1.4 BatchMatmul: 513.9 us

• tf-r1.4 BatchMatmul + Transpose (separate): 541.9 us

• TVM BatchMatmul: 37.62 us

• TVM BatchMatmul + Transpose (fused): 38.39 us

内核融合优化带来了 1.7 倍的加速。

集成 TensorFlow

在我们的工作负载中，batch 矩阵相乘的输入形状是有限的，易于提前枚举。有了这些预定义的形状，我们可以提前生成高度优化的 CUDA 内核（固定形状的计算可以带来最佳优化潜能）。同时，还将生成一个适合大多数形状的通用 batch 矩阵相乘内核，为没有提前生成内核的形状提供回退机制。

我们将生成的针对特定形状的高效内核和回退机制集成到 Tensorflow 中。我们开发了一些融合操作，例如 BatchMatMulTranspose 或 BatchMatMulAdd——使用 TVM runtime API 为确定输入形状启动特定生成的内核或调用回退内核。

通过执行图优化 pass，可以利用融合操作自动替换原始batch matmul + add/transpose。同时，通过结合更激进的图优化 pass，我们尝试利用 TVM 为长尾操作模式生成更高效的融合内核，以进一步提升端到端性能。

总结

在阿里巴巴，我们发现 TVM 是非常有效的开发高性能 GPU 内核的工具，可以满足我们的内部需求。

在本博客中，我们以 Transformer 模型为例，说明了我们利用 TVM 的优化策略。

首先，我们通过第一原理分析确定了 Transformer 模型的关键问题。然后，我们使用 TVM 生成高度优化的 CUDA 内核来取代 cuBLAS 版本（此时达到 13 倍的加速）。

接下来，利用 TVM 的内核融合机制来融合 batch 矩阵相乘的前/后操作，以带来进一步的性能改进（性能提升 1.7 倍）。端到端的性能改善达到 1.4 倍。基于这些生成的内核，我们开发了一个图优化 pass，自动用 TVM 融合内核替换掉原有的计算模式，确保优化对终端用户透明。

最后，所有这些优化都以松散耦合的方式集成到 TensorFlow 中，这展示了将 TVM 与不同深度学习框架集成的潜在方式。

目前我们还有一项正在进行的工作——将 TVM 整合为 TensorFlow 的 codegen 后端。我们希望将来与社群分享更多成果。

via： http://www.tvmlang.org

雷锋网 (公众号：雷锋网) AI 研习社编译整理。

雷锋网版权文章，未经授权禁止转载。详情见。