TVM之Schedule

admin管理员组

文章数量:1531435

TVM之Schedule

• • 1. schedule
  • 2. stage
  • 3. methods for stage
  • • 3.1 split
    • 3.2 tile
    • 3.3 fuse
    • 3.4 reorder
    • 3.5 bind
    • 3.6 compute_at
    • 3.7 compute_inline
    • 3.8 compute_root

最近在阅读TVM官网的文档，在阅读的过程中将keywords以及相应理解记录了下来，有理解不对的地方，欢迎批评指正~~
update：2019/07/07 凌晨更新，欢迎提意见，欢迎来喷 ?

官网文档地址：Schedule Primitives in TVM

1. schedule

There often exist several methods to compute the same result, however, different methods will result in different locality and performance. So TVM asks user to provide how to execute the computation called Schedule. [摘自官方文档]

在描述了算子逻辑之后，用户需要告诉TVM如何具体地计算。通常情况下，有多种方式可以计算得到相同的结果，但是执行起来的效率可能相差很大，例如计算从1累加到100的结果，就有如下两种方式。因此，使用方式1还是方式2，这个过程就是schedule。

int sum = 0;

//第一种方式：循环累加
for(int i=1; i<101, i++)
{
   
    sum = sum + i;
}

//第二种方式：公式计算
sum = (1 + 100) * 50;

如下的例子中，计算shape均为(m,n)的两个tensor的乘积：

A = tvm.placeholder((m, n), name='A')
B = tvm.placeholder((m, n), name='B')
C = tvm.compute((m, n), lambda i, j: A[i, j] * B[i, j], name='C')

s = tvm.create_schedule([C.op])
# lower will transform the computation from definition to the real
# callable function. With argument `simple_mode=True`, it will
# return you a readable C like statement, we use it here to print the
# schedule result.
print(tvm.lower(s, [A, B, C], simple_mode=True))

使用tvm.create_schedule会产生一个默认的schedule。该默认的schedule如下所示，通过两层for循环计算C每个元素的积，该schedule未经过优化操作。

produce C {
   
  for (i, 0, m) {
   
    for (j, 0, n) {
   
      C[((i*n) + j)] = (A[((i*n) + j)]*B[((i*n) + j)])
    }
  }
}

---

2. stage

One schedule is composed by multiple stages, and one Stage represents schedule for one operation. We provide various methods to schedule every stage. [摘自官方文档]

此时引入了Stage的概念，一个完整的schedule包含多个stage，而每个stage是对一个操作的schedule。这里stage姑且理解为sub schedule。

TVM中有很多的方法可以用于每个stage。可以理解为，这些方法可用于优化默认的schedule。

---

3. methods for stage

3.1 split

split can split a specified axis into two axises by factor. [摘自官方文档]

该函数可以将tensor按某个指定的轴进行切分，如下的例子中计算tensor A 乘以标量2的结果，默认的schedule为一个for(int i=0; i<m-1; i++)的一个for循环，通过如下的例子中的split函数，可以将一个for循环切割成两个嵌套的循环。

A = tvm.placeholder((m,), name='A')
B = tvm.compute((m,), lambda i: A[i]*2, name='B')

s = tvm.create_schedule(B.op)
xo, xi = s[B].split(B.op.axis[0], factor=32)
print(tvm.lower(s, [A, B], simple_mode=True))

split之后的schedule如下所示，假设m=35，则A的shape为(35,)，且A的最后一个值为A[34]。split之后的计算过程变为两层for循环，其中i.outer的取值范围为[0,1]，当i.outer为0时，i.inner的取值范围为[0,31]。当i.outer为1时，i.inner的取值范围为[0,2]。因此最后一个值的计算过程为：B[32+2] = A[32+2] * 2。

produce B {
   
  for (i.outer, 0, ((m + 31)/32)) {
   
    for (i.inner, 0, 32) {
   
      if (likely(((i.outer*32) < (m - i.inner)))) {
   
        B[((i.outer*32) + i.inner)] = (A[((i.outer*32) + i.inner)]*2.000000f)
      }
    }
  }
}

3.2 tile

tile help you execute the computation tile by tile over two axises. [摘自官方文档]

该函数类似于split，区别是tile可以对2层for循环进行split操作。如下代码所示，其功能为输出一个与输入tensor A内容相同的tensor B。默认的schedule为for(int i=0; i<m-1; i++)与for(int j=0; j<n-1; j++)的两层for循环。通过tile可以指定0号轴按10进行切分，指定1号轴按5进行切分。

A = tvm.placeholder((m, n

本文标签： TVMschedule