TVM调度原语官方DOC
#!/usr/bin/env python
# coding: utf-8
# 通常有几种方法来计算相同的结果，但是不同的方法会导致不同的局部性和性能。所以TVM要求用户提供如何执行称为Schedule的计算。
# Schedule是一组转换计算的集合，它转换程序中的计算循环。
# In[1]:
#兼容python2.x
from __future__ import print_function #在python2.x中也要像python3.x那样使用print（）
from __future__ import absolute_import #在python2.x中使用绝对路径
# In[2]:
import tvm
from tvm import te
import numpy as np
# In[3]:
#声明两个var变量
n = te.var("n") #使用指定的名称和dtype创建一个新变量,dtype默认为‘int32’
m = te.var("m")
# In[4]:
#声明一个element-wise(相对位置相乘)乘法
A = te.placeholder((m, n), name="A") #构造一个空张量对象
B = te.placeholder((m, n), name="B")
C = te.compute((m, n), lambda i, j: A[i, j] * B[i, j], name="C") #声明运算
s = te.create_schedule([C.op]) #创建调度
print(tvm.lower(s, [A, B, C], simple_mode=True)) #lower将把计算从定义转换为实际的可调用函数
# # split
# In[5]:
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) #split用factor参数把一个指定的轴分解为两个轴
print(tvm.lower(s, [A, B], simple_mode=True))
# In[6]:
A = te.placeholder((m,), name="A")
B = te.compute((m,), lambda i:A[i], name='B')
s = te.create_schedule(B.op)
bx, tx = s[B].split(B.op.axis[0], nparts=32) #使用nparts与使用factor正好相反
print(tvm.lower(s, [A, B], simple_mode = True))
# # tile
# In[7]:
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，把数据进行平铺，例如这里的二维数据，相当于沿着两个轴，把数据切割成了一个一个小块
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
# In[8]:
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)
fused = s[B].fuse(xi, yi) #把xi轴和yi轴融合成了一个轴
print(tvm.lower(s, [A, B], simple_mode=True))
# # reorder
# In[9]:
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)
s[B].reorder(xi, yo, xo, yi) #将轴重新排序,即改变计算过程中的循环嵌套顺序
print(tvm.lower(s, [A, B], simple_mode=True))
# # bind&thread_axis
# In[10]:
A = te.placeholder((n,), name="A")
B = te.compute(A.shape, lambda i: A[i] * 2, name="B")
s = te.create_schedule(B.op)
bx, tx = s[B].split(B.op.axis[0], factor=64)
s[B].bind(bx, te.thread_axis("blockIdx.x")) #thread_axis:创建迭代变量来表示线程索引.bind:将指定的轴与线程绑定,通常用于gpu编程
s[B].bind(tx, te.thread_axis("threadIdx.x"))
print(tvm.lower(s,[A, B], simple_mode=True))
# # compute_at
# In[11]:
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)
print(tvm.lower(s, [A, B, C], simple_mode=True)) #对于由多个运算符组成的调度,TVM默认情况下会在root分别计算张量.
# In[13]:
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]) #compute_at将B的计算挪到了C的计算的第一个轴里
print(tvm.lower(s, [A, B, C], simple_mode=True))
# # compute_inline
# In[16]:
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() #compute_inline()将一个阶段标记为內联,然后计算体将被展开并插入到需要它的地方.
print(tvm.lower(s,[A, B, C], simple_mode=True))
# # copute_root
# In[18]:
A = te.placeholder((m,), name = "A")
B = te.compute((m,), lambda i: A[i] + 1)
C = te.compute((m,), lambda i: B[i] * 2)
s = te.create_schedule(C.op)
s[B].compute_at(s[C], C.op.axis[0])
s[B].compute_root() #将B的计算阶段挪回root
print(tvm.lower(s, [A, B, C], simple_mode = True))
# In[ ]: