TVM Compiler中文教程：TVM如何优化CPU GEMM(矩阵乘法)

TVM如何优化CPU GEMM(矩阵乘法)

TVM提供抽象接口，允许用户分别描述算法和算法的实施组织（所谓的调度Schedule）。通常，写高性能调度的算法时，会破坏算法的可读性和模块性。此外，尝试各种看似有用的调度是非常耗费人力的。在TVM帮助下，我们能高效地尝试这些调度，来提高计算性能。

在本教程中，我们将演示如何使用TVM优化方阵乘法，并通过简单地添加18行额外代码，实现比基线版本快200倍的优化版本。

在CPU上执行密集计算有两个优化重点：

1. 提高内存访问的cache命中率。复杂的数值计算和hot-spot存储器访问都可以从高cache命中率中加速。这要求我们将原始内存访问模式转换为符合cache策略的模式。
2. SIMD（单指令多数据），也成为向量化处理单元。每次都会处理一小批数据，而不是单个网格。这要求我们以统一模式在循环体中转变数据访问模式，以便LLVM后端可以将其降低到SIMD。

实际上，本教程中使用的所有方法都是此repo中提到的技巧的子集。其中一些已经被TVM抽象并自动应用，但由于TVM限制，其中一些不能简单地应用。

下面提到的所有实验结果都是在配备Intel i7-4770HQ CPU的2015款15英寸MacBook上执行的。对于所有x86 CPU，高速缓存行大小应为64字节。

• 注：代码里面的测试时间是个人在E5-2620 v4测得

准备和基线

在本教程中，我们将演示如何使用TVM来优化矩阵乘法。在实际演示之前，我们首先定义这些变量。然后我们编写一个基线实现，这是在TVM中编写矩阵乘法的最简单方法。

import tvm
import numpy
import timeit

#矩阵大小(M,K)x(K,N),可以自由尝试不同的形状，有时TVM优化优于MKL的numpy。
M = 1024
K = 1024
N = 1024
#TVM中默认张量类型
dtype = "float32"
#使用Intel AVX2（高级矢量扩展）ISA进行SIMD
#获得最佳新能要修改下一行为'llvm -mcpu=core-avx2'，或者指定你使用的其他CPU类型
#实测指定CPU以后，Opt6版本可以达到略高于MKL numpy的性能。
target = 'llvm'
ctx = tvm.context(target, 0)

# Random generated tensor for testing
a = tvm.nd.array(numpy.random.rand(M, K).astype(dtype), ctx)
b = tvm.nd.array(numpy.random.rand(K, N).astype(dtype), ctx)

#numpy测试Numpy running time: 0.004753ms
np_repeat = 100
np_runing_time = timeit.timeit(setup='import numpy\n'
                                     'M = ' + str(M) + '\n'
                                     'K = ' + str(K) + '\n'
                                     'N = ' + str(N) + '\n'
                                     'dtype = "float32"\n'
                                     'a = numpy.random.rand(M, K).astype(dtype)\n'
                                     'b = numpy.random.rand(K, N).astype(dtype)\n',
                               stmt='answer = numpy.dot(a, b)',
                               number=np_repeat)
print("Numpy running time: %f" % (np_runing_time / np_repeat))

#基准测试：Baseline: 2.174880ms
#算法
k = tvm.reduce_axis((0, K), 'k')
A = tvm.placeholder((M, K), name='A')
B = tvm.placeholder((K, N), name='B')
C = tvm.compute(
           (M, N),
           lambda x, y: tvm.sum(A[x, k] * B[k, y], axis=k),
           name='C')
#默认调度
c = tvm.nd.array(numpy.zeros((M, N), dtype=dtype), ctx)
func(a, b, c)
tvm.testing.assert_allclose(c.asnumpy(), answer, rtol=1e-5)

evaluator = func.time_evaluator(func.entry_name, ctx, number=1)
print('Baseline: %f' % evaluator(a, b, c).mean)

输出：

Numpy running time: 0.036594
Baseline: 2.533367

在TVM中，我们可以随时检查较低级别的IR以调试或优化我们的调度。以下是使用我们的基线计划生成的IR。

print(tvm.lower(s, [A, B, C], simple_mode=True))

输出：

produce C {
   
  for (x, 0, 1024) {
   
    for (y, 0, 1024) {
   
      C[((x*1024) + y)] = 0.000000f
      for (k, 0, 1024) {
   
        C[((x*1024) + y)] = (C[((x*1024) + y)] + (A[((x*1024) + k)]*B[((k*1024) + y)]))
      }
    }
  }
}

Opt1：分块

分块（数据块将逐块计算）是一个增强cache命中率的重要技巧。块内的存储器访问是一个具有高内存局部性的小邻域。在本教程中，我们选择32作为分块因子。因此该块将填充32 * 32 * sizeof（float），它占总大小为32KB的缓存中的4KB（L1数据缓存）。

#Opt1调度方法：
bn = 32
#通过循环平铺tile来分块
xo, yo, xi, yi = s[C].tile(C.op.axis[0], C.op.axis[1], bn, bn)#32x32的块，返回内外循环索引
k, = s[C].op.reduce_axis
ko, ki = s[C].split(k, factor=4)
#ko,ki移到内外块循环之间
s[C].reorder(xo, yo, ko, ki, xi, yi)

func = tvm.build(s, [A, B, C], target=target, name='mmult')
assert func

c = tvm.nd.array(numpy.zeros((M, N), dtype = dtype), ctx)
func(a, b, c)
tvm.testing.assert_allclose(c.asnumpy(), answer, rtol=1e-5)
#通过简单地平铺循环32x32，并在外分块循环出加入ko，ki循环，我们可以看到与基线相比，有很大加速。
#Opt1: 0.382091ms
evaluator = func.time_evaluator(func.entry_name, ctx, number=10)
print('Opt1: %f' % evaluator(a, b, c).mean)

输出：

Opt1: 0.296531

这是在分块后生成的IR。

print(tvm.lower(s, [A, B, C], simple_mode=True))

输出：

produce C {
   
  for (x.outer, 0, 32) {
   
    for (y.outer, 0, 32) {
   
      for (x.inner.init, 0, 32) {
   
        for (y.inner.init, 0, 32) {
   
          C[((((((x.outer*32) + x.inner.init)*32) + y.outer)*32) + y.inner.init)] = 0.000000f
        }
      }
      for (k.outer, 0, 256) {
   
        for (k.inner, 0, 4) {
   
          for (x.inner, 0, 32) {
   
            for (y.inner, 0, 32) {
   
              C[((((((x.outer*32) + x.inner)*32) + y.outer)*32) + y.inner)] = (C[((((((x.outer*32) + x.inner)*32) + y.outer)*32) + y.inner)] + (A[((((((x.outer*32) + x.inner)*256) + k.outer)