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[ ]: