TVMCompiler中⽂教程：TVM如何优化CPUGEMM（矩阵乘法）
⽂章⽬录
TVM如何优化CPU GEMM(矩阵乘法)
TVM提供抽象接⼝，允许⽤户分别描述算法和算法的实施组织（所谓的调度Schedule）。通常，写⾼性能调度的算法时，会破坏算法的可读性和模块性。此外，尝试各种看似有⽤的调度是⾮常耗费⼈⼒的。在TVM帮助下，我们能⾼效地尝试这些调度，来提⾼计算性能。
在本教程中，我们将演⽰如何使⽤TVM优化⽅阵乘法，并通过简单地添加18⾏额外代码，实现⽐基线版本快200倍的优化版本。
在CPU上执⾏密集计算有两个优化重点：
1. 提⾼内存访问的cache命中率。复杂的数值计算和hot-spot存储器访问都可以从⾼cache命中率中加速。这要求我们将原始内存访问
模式转换为符合cache策略的模式。
2. SIMD（单指令多数据），也成为向量化处理单元。每次都会处理⼀⼩批数据，⽽不是单个⽹格。这要求我们以统⼀模式在循环体中转
变数据访问模式，以便LLVM后端可以将其降低到SIMD。
风车做法
实际上，本教程中使⽤的所有⽅法都是此中提到的技巧的⼦集。其中⼀些已经被TVM抽象并⾃动应⽤，但由于TVM限制，其中⼀些不能简单地应⽤。
下⾯提到的所有实验结果都是在配备Intel i7-4770HQ CPU的2015款15英⼨MacBook上执⾏的。对于所有x86 CPU，⾼速缓存⾏⼤⼩应为64字节。
注：代码⾥⾯的测试时间是个⼈在E5-2620 v4测得
准备和基线
在本教程中，我们将演⽰如何使⽤TVM来优化矩阵乘法。在实际演⽰之前，我们⾸先定义这些变量。然后我们编写⼀个基线实现，这是在TVM中编写矩阵乘法的最简单⽅法。
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 = t(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(tup='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))
#基准测试：Baline: 2.174880ms
#算法
k = duce_axis((0, K),'k')
A = tvm.placeholder((M, K), name='A')
B = tvm.placeholder((K, N), name='B')
C = pute(
(M, N),
lambda x, y: tvm.sum(A[x, k]* B[k, y], axis=k),
name='C')
#默认调度
c = tvm.nd.s((M, N), dtype=dtype), ctx)
func(a, b, c)
纬线的特点
evaluator = func.time__name, ctx, number=1)
print('Baline: %f'% evaluator(a, b, c).mean)
输出：
Numpy running time: 0.036594
Baline: 2.533367
在TVM中，我们可以随时检查较低级别的IR以调试或优化我们的调度。以下是使⽤我们的基线计划⽣成的IR。print(tvm.lower(s,[A, B, C], simple_mode=True))
输出：
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')
分手告别的话asrt func
c = tvm.nd.s((M, N), dtype = dtype), ctx)
func(a, b, c)
#通过简单地平铺循环32x32，并在外分块循环出加⼊ko，ki循环，我们可以看到与基线相⽐，有很⼤加速。
#Opt1: 0.382091ms
evaluator = func.time__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))
输出：
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.in ner)*256)+ k.outer)*4)+ k.inner)]*B[((((((k.outer*4)+ k.inner)*32)+ y.outer)*32)+ y.inner)]))
}
}
}
}
}
}
电子元器件知识}
Opt2：向量化
向量化是另⼀个重要的技巧。当内存访问模式⼀致时，编译器可以检测到这种模式并将连续内存传递给向量处理器。在TVM中，我们可以使⽤vectorize接⼝来暗⽰编译器这种模式，这样我们就可以⼤⼤加速它。 在本教程中，我们选择对内部循环⾏数据进⾏⽮量化，因为它是cache友好的。
s = ate_schedule(C.op)
xo, yo, xi, yi = s[C].tile(C.op.axis[0], C.op.axis[1], bn, bn)
k,= s[C].op.reduce_axis
ko, ki = s[C].split(k, factor=4)
s[C].reorder(xo, yo, ko, ki, xi, yi)
# 与Opt1唯⼀区别在：向量化
s[C].vectorize(yi)
func = tvm.build(s,[A, B, C], target=target, name='mmult')
asrt func
c = tvm.nd.s((M, N), dtype = dtype), ctx)
func(a, b, c)感激之礼
#Opt2: 0.356233ms
evaluator = func.time__name, ctx, number=10)
print('Opt2: %f'% evaluator(a, b, c).mean)
输出：
Opt2: 0.312343
这是向量化后⽣产的IR。
print(tvm.lower(s,[A, B, C], simple_mode=True))
输出：
C[ramp((((((x.outer*32)+ x.inner.init)*32)+ y.outer)*32),1,32)]=x32(0.000000f)
}
for(k.outer,0,256){
for(k.inner,0,4){
for(x.inner,0,32){
C[ramp((((((x.outer*32)+ x.inner)*32)+ y.outer)*32),1,32)]=(C[ramp((((((x.outer*32)+ x.inner)*32)+ y.outer)*32),1,32)]+(x32(A[((((((x.outer*32 )+ x.inner)*256)+ k.outer)*4)+ k.inner)])*B[ramp((((((k.outer*4)+ k.inner)*32)+ y.outer)*32),1,32)]))
}
}
}
}
天津市一日游最好地方
}
}
Opt3：循环排布permute
如果我们看⼀下上⾯的IR，我们可以看到内部循环⾏数据被向量化，B被转换成PackedB。PackedB
的遍历现在是顺序的。因此，我们将查看A的访问模式。在当前调度中，A逐列访问，这是cache不友好的。如果我们改变ki和内轴xi的嵌套循环次序，则A矩阵的访问模式更加cache友好。
s = ate_schedule(C.op)
xo, yo, xi, yi = s[C].tile(C.op.axis[0], C.op.axis[1], bn, bn)
k,= s[C].op.reduce_axis
ko, ki = s[C].split(k, factor=4)
# re-ordering，改变xi和ki的循环位置
s[C].reorder(xo, yo, ko, xi, ki, yi)
s[C].vectorize(yi)
func = tvm.build(s,[A, B, C], target=target, name='mmult')
asrt func
c = tvm.nd.s((M, N), dtype = dtype), ctx)
func(a, b, c)
#Opt3: 0.098853
evaluator = func.time__name, ctx, number=10)
print('Opt3: %f'% evaluator(a, b, c).mean)
输出：
Opt3: 0.141706
这是循环重排后⽣成的IR。
print(tvm.lower(s,[A, B, C], simple_mode=True))
输出：
