第一件事情當(dāng)然是安裝苛坚，官方給了教程。
- https://tvm.apache.org/docs/install/index.html撒踪，配置文件中要啟用cuda,cublas,llvm
功能
深度學(xué)習(xí)模型的編譯
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,這是增加循環(huán)嵌套層數(shù)同時不改變循環(huán)體,從并行計(jì)算的角度講衩婚，分割后的外循環(huán)的每一次迭代都可能獨(dú)立執(zhí)行，這為并行計(jì)算加速創(chuàng)造了條件助赞。不過如果原有循環(huán)體中的數(shù)據(jù)關(guān)聯(lián)程度很大买羞，導(dǎo)致外循環(huán)中的不同次迭代產(chǎn)生關(guān)聯(lián)，此時就不能夠按照這種方法并行計(jì)算雹食。下面是樣例以及修改后的對比畜普。
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))
分割后的結(jié)果，相當(dāng)于增加了一層循環(huán)嵌套:
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,增加二維數(shù)組的嵌套群叶，可以同時給第0個軸和第1個軸增加循環(huán)嵌套吃挑。
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運(yùn)算是分割split運(yùn)算的逆運(yùn)算，將兩個循環(huán)合并到一層,
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街立，重排序是指按照順序?qū)⑤S重新排序舶衬，做這個操作的目的可以增大cache的命中率。二維數(shù)組可以把它看做是由若干行所構(gòu)成的結(jié)構(gòu)赎离。循環(huán)體的訪存可以看作3次讀取逛犹，一次寫入。這種寫法保證了按k增加的二維數(shù)組的訪存局部性A[i,k]友好梁剔，但是對于C[i,j]和B[k,j]都是不友好的圾浅。我們于是設(shè)想把k提出來，對j進(jìn)行優(yōu)化憾朴。
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的計(jì)算移入C的計(jì)算的第一軸狸捕。這樣我們可以將兩個循環(huán)合并成為一個循環(huán)。
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() # 可以使用這個函數(shù)將其再移動回去众雷。
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,使用內(nèi)聯(lián)的方法灸拍。
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)
}
}
