第一件事情当然是安装，官方给了教程。
- https://tvm.apache.org/docs/install/index.html，配置文件中要启用cuda,cublas,llvm
功能
深度学习模型的编译
https://zhuanlan.zhihu.com/p/335374653
Expressions for Operators
参考链接
https://tvm.apache.org/docs/tutorials/get_started/tensor_expr_get_started.html#sphx-glr-tutorials-get-started-tensor-expr-get-started-py 及https://zhuanlan.zhihu.com/p/267671270
Describe the Computation
n = te.var("n")
A = te.placeholder((n,), name="A")
B = te.placeholder((n,), name="B")
C = te.compute(A.shape, lambda i: A[i] + B[i], name="C")
s = te.create_schedule([C.op])
print(tvm.lower(s, [A, B, C], simple_mode=True)) # 可以查看schedule result
这个和编写tf代码类似。具体如何使用参照链
schedule的几种方法。
- Split,这是增加循环嵌套层数同时不改变循环体,从并行计算的角度讲，分割后的外循环的每一次迭代都可能独立执行，这为并行计算加速创造了条件。不过如果原有循环体中的数据关联程度很大，导致外循环中的不同次迭代产生关联，此时就不能够按照这种方法并行计算。下面是样例以及修改后的对比。
A = te.placeholder((m,), name="A")
B = te.compute((m,), lambda i: A[i] * 2, name="B")
s = te.create_schedule(B.op)
xo, xi = s[B].split(B.op.axis[0], factor=32) # 此处可以使用 nparts=32 分割相反的轴
print(tvm.lower(s, [A, B], simple_mode=True))
分割后的结果，相当于增加了一层循环嵌套:
primfn(A_1: handle, B_1: handle) -> ()
attr = {"global_symbol": "main", "tir.noalias": True}
buffers = {B: Buffer(B_2: Pointer(float32), float32, [m: int32], [stride: int32], type="auto"),
A: Buffer(A_2: Pointer(float32), float32, [m], [stride_1: int32], type="auto")}
buffer_map = {A_1: A, B_1: B} {
for (i.outer: int32, 0, floordiv((m + 31), 32)) {
for (i.inner: int32, 0, 32) {
if @tir.likely((((i.outer*32) + i.inner) < m), dtype=bool) {
B_2[(((i.outer*32) + i.inner)*stride)] = ((float32*)A_2[(((i.outer*32) + i.inner)*stride_1)]*2f32)
}
}
}
}
- Tile,增加二维数组的嵌套，可以同时给第0个轴和第1个轴增加循环嵌套。
A = te.placeholder((m, n), name="A")
B = te.compute((m, n), lambda i, j: A[i, j], name="B")
s = te.create_schedule(B.op)
xo, yo, xi, yi = s[B].tile(B.op.axis[0], B.op.axis[1], x_factor=10, y_factor=5)
print(tvm.lower(s, [A, B], simple_mode=True))
- Fuse,合并fuse运算是分割split运算的逆运算，将两个循环合并到一层,
B = te.compute((m, n), lambda i, j: A[i, j], name="B")
s = te.create_schedule(B.op)
# tile to four axises first: (i.outer, j.outer, i.inner, j.inner)
xo, yo, xi, yi = s[B].tile(B.op.axis[0], B.op.axis[1], x_factor=10, y_factor=5)
# then fuse (i.inner, j.inner) into one axis: (i.inner.j.inner.fused)
fused = s[B].fuse(xi, yi)
print(tvm.lower(s, [A, B], simple_mode=True))
- Reorder，重排序是指按照顺序将轴重新排序，做这个操作的目的可以增大cache的命中率。二维数组可以把它看做是由若干行所构成的结构。循环体的访存可以看作3次读取，一次写入。这种写法保证了按k增加的二维数组的访存局部性A[i,k]友好，但是对于C[i,j]和B[k,j]都是不友好的。我们于是设想把k提出来，对j进行优化。
for i in range(1024):
for j in range(1024):
for k in range(32):
C[i,j] += A[i,k] * B[k,j]
# reorder后的操作
for k in range(32):
for i in range(1024):
for j in range(1024):
C[i,j] += A[i,k] * B[k,j]
TVM中使用方法:
A = te.placeholder((m, n), name="A")
B = te.compute((m, n), lambda i, j: A[i, j], name="B")
s = te.create_schedule(B.op)
# tile to four axises first: (i.outer, j.outer, i.inner, j.inner)
xo, yo, xi, yi = s[B].tile(B.op.axis[0], B.op.axis[1], x_factor=10, y_factor=5)
# then reorder the axises: (i.inner, j.outer, i.outer, j.inner)
s[B].reorder(xi, yo, xo, yi)
print(tvm.lower(s, [A, B], simple_mode=True))
- 其余的一些方法:
B = te.compute((m,), lambda i: A[i] + 1, name="B")
C = te.compute((m,), lambda i: B[i] * 2, name="C")
s = te.create_schedule(C.op)
print(tvm.lower(s, [A, B, C], simple_mode=True))
primfn(A_1: handle, B_1: handle, C_1: handle) -> ()
attr = {"global_symbol": "main", "tir.noalias": True}
buffers = {C: Buffer(C_2: Pointer(float32), float32, [m: int32], [stride: int32], type="auto"),
B: Buffer(B_2: Pointer(float32), float32, [m], [stride_1: int32], type="auto"),
A: Buffer(A_2: Pointer(float32), float32, [m], [stride_2: int32], type="auto")}
buffer_map = {A_1: A, B_1: B, C_1: C} {
for (i: int32, 0, m) {
B_2[(i*stride_1)] = ((float32*)A_2[(i*stride_2)] + 1f32)
}
for (i_1: int32, 0, m) {
C_2[(i_1*stride)] = ((float32*)B_2[(i_1*stride_1)]*2f32)
}
}
我们将使用compute_at将B的计算移入C的计算的第一轴。这样我们可以将两个循环合并成为一个循环。
A = te.placeholder((m,), name="A")
B = te.compute((m,), lambda i: A[i] + 1, name="B")
C = te.compute((m,), lambda i: B[i] * 2, name="C")
s = te.create_schedule(C.op)
s[B].compute_at(s[C], C.op.axis[0])
s[B].compute_root() # 可以使用这个函数将其再移动回去。
print(tvm.lower(s, [A, B, C], simple_mode=True))
primfn(A_1: handle, B_1: handle, C_1: handle) -> ()
attr = {"global_symbol": "main", "tir.noalias": True}
buffers = {B: Buffer(B_2: Pointer(float32), float32, [m: int32], [stride: int32], type="auto"),
C: Buffer(C_2: Pointer(float32), float32, [m], [stride_1: int32], type="auto"),
A: Buffer(A_2: Pointer(float32), float32, [m], [stride_2: int32], type="auto")}
buffer_map = {A_1: A, B_1: B, C_1: C} {
for (i: int32, 0, m) {
B_2[(i*stride)] = ((float32*)A_2[(i*stride_2)] + 1f32)
C_2[(i*stride_1)] = ((float32*)B_2[(i*stride)]*2f32)
}
}
compute_inline,使用内联的方法。
A = te.placeholder((m,), name="A")
B = te.compute((m,), lambda i: A[i] + 1, name="B")
C = te.compute((m,), lambda i: B[i] * 2, name="C")
s = te.create_schedule(C.op)
s[B].compute_inline()
print(tvm.lower(s, [A, B, C], simple_mode=True))
primfn(A_1: handle, B_1: handle, C_1: handle) -> ()
attr = {"global_symbol": "main", "tir.noalias": True}
buffers = {C: Buffer(C_2: Pointer(float32), float32, [m: int32], [stride: int32], type="auto"),
B: Buffer(B_2: Pointer(float32), float32, [m], [stride_1: int32], type="auto"),
A: Buffer(A_2: Pointer(float32), float32, [m], [stride_2: int32], type="auto")}
buffer_map = {A_1: A, B_1: B, C_1: C} {
for (i: int32, 0, m) {
C_2[(i*stride)] = (((float32*)A_2[(i*stride_2)] + 1f32)*2f32)
}
}
