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), name='A')
B = tvm.compute((m, n), lambda i, j: A[i, j], name='B')

s = tvm.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))

使用tile进行schedule后，生成的代码如下。此时内层出现了分别为10与5的两层for循环。看官可能要问为何这么操作，至于这个我是知道的，但是我不告诉你 ?

produce B {
  for (i.outer, 0, ((m + 9)/10)) {
    for (j.outer, 0, ((n + 4)/5)) {
      for (i.inner, 0, 10) {
        for (j.inner, 0, 5) {
          if (likely((((i.outer*10) + i.inner) < m))) {
            if (likely((((j.outer*5) + j.inner) < n))) {
              if (likely((((i.outer*10) + i.inner) < m))) {
                if (likely((((j.outer*5) + j.inner) < n))) {
                  B[(((j.outer*5) + (((i.outer*10) + i.inner)*n)) + j.inner)] = A[(((j.outer*5) + (((i.outer*10) + i.inner)*n)) + j.inner)]
                }
              }
            }
          }
        }
      }
    }
  }
}

3.3 fuse

fuse can fuse two consecutive axises of one computation.[摘自官方文档]

一句话概括fuse的作用：将连续2根轴（for循环）融合为1根轴（for循环）。

基于3.2 tile中的例子，2层for循环变为4层for循环，最内2层的for循环为[0 ~ 10)与[0 ~ 5)，那么为了能达到某些目标，这2层for循环可以合并为一层for循环，可以使用[0 ~ 50)这个for循环来替换。

如下代码在3.2 tile的基础上增加了fused = s[B].fuse(xi, yi)，其作用是融合最后2个for循环。所以在输出的ir中就只有3个for循环。

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

s = tvm.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))

produce B {
  for (i.outer, 0, ((m + 9)/10)) {
    for (j.outer, 0, ((n + 4)/5)) {
      for (i.inner.j.inner.fused, 0, 50) {
        if (likely((((i.outer*10) + (i.inner.j.inner.fused/5)) < m))) {
          if (likely((((j.outer*5) + (i.inner.j.inner.fused % 5)) < n))) {
            if (likely((((i.outer*10) + (i.inner.j.inner.fused/5)) < m))) {
              if (likely((((j.outer*5) + (i.inner.j.inner.fused % 5)) < n))) {
                B[(((j.outer*5) + (((i.outer*10) + (i.inner.j.inner.fused/5))*n)) + (i.inner.j.inner.fused % 5))] = A[(((j.outer*5) + (((i.outer*10) + (i.inner.j.inner.fused/5))*n)) + (i.inner.j.inner.fused % 5))]
              }
            }
          }
        }
      }
    }
  }
}

3.4 reorder

reorder can reorder the axises in the specified order.[摘自官方文档]

reorder接口可以使得多层for循环按指定的次序重新排列，以3.2 tile输出的ir为例，输出的4层的for循环分别为：

  1. 0 ~ ((m + 9)/10)
  2. 0 ~ ((n + 4)/5)
  3. 0 ~ 10
  4. 0 ~ 5

基于某个原因，需要得到的for循环的次序如下所示：

  1. 0 ~ 10
  2. 0 ~ ((n + 4)/5)
  3. 0 ~ ((m + 9)/10)
  4. 0 ~ 5

此时就需要reorder显威力了，在原有代码的基础上，使用s[B].reorder(xi, yo, xo, yi)调整for循环的顺序。

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

s = tvm.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))

使得原来的第1层与第3层互换位置，从而得到新的for循环代码。看官可将如下的ir与3.2 tile中的输出ir进行对比。在实际的开发过程中，该招数还是很有用滴。。。。。

produce B {
  for (i.inner, 0, 10) {
    for (j.outer, 0, ((n + 4)/5)) {
      for (i.outer, 0, ((m + 9)/10)) {
        for (j.inner, 0, 5) {
          if (likely((((i.outer*10) + i.inner) < m))) {
            if (likely((((j.outer*5) + j.inner) < n))) {
              if (likely((((i.outer*10) + i.inner) < m))) {
                if (likely((((j.outer*5) + j.inner) < n))) {
                  B[(((j.outer*5) + (((i.outer*10) + i.inner)*n)) + j.inner)] = A[(((j.outer*5) + (((i.outer*10) + i.inner)*n)) + j.inner)]
                }
              }
            }
          }
        }
      }
    }
  }
}

3.5 bind

bind can bind a specified axis with a thread axis, often used in gpu programming.[摘自官方文档]

这个bind有多牛捏，介么跟你形容吧，并行计算侬晓得伐，用bind才能进行万马奔腾的计算。看，上上上上面的3.1中的例子，虽然1层for循环被split成了2层for循环，但是对于底层硬件来说，并没有啥区别。其执行起来仍然是单核心的串行计算，那么对于像GPU这种加速芯片来说，完全就体现不出并行计算的优势。如果需要充分的利用并行的性能，那么就需要将计算任务合理地分配到每个计算单元中。

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

s = tvm.create_schedule(B.op)
bx, tx = s[B].split(B.op.axis[0], factor=64)
s[B].bind(bx, tvm.thread_axis("blockIdx.x"))
s[B].bind(tx, tvm.thread_axis("threadIdx.x"))
print(tvm.lower(s, [A, B], simple_mode=True))

以下面的代码为例，基于3.1中的例子，此时得到的ir为两层for循环。假设n的值为128，那么两层for循环分别为[0,2)与[0,64)。将[0,2)这根轴绑定blockIdx.x，那么就可以理解为，将该计算任务分配到2个核心中进行计算，每个核心的计算量是[0,64)这个for循环的计算量。

基于核心的粒度分配好任务量后，在每个核心中，处理的是[0,64)这个for循环，将该轴绑定到threadIdx.x，那么就可以理解为每个核开启多线程，每个线程处理的内容为[0,64)这个for循环中的任务（本例子中任务就是一个乘以2的操作）。

produce B {
  // attr [iter_var(blockIdx.x, , blockIdx.x)] thread_extent = ((n + 63)/64)
  // attr [iter_var(threadIdx.x, , threadIdx.x)] thread_extent = 64
  if (likely((((blockIdx.x*64) + threadIdx.x) < n))) {
    if (likely((((blockIdx.x*64) + threadIdx.x) < n))) {
      B[((blockIdx.x*64) + threadIdx.x)] = (A[((blockIdx.x*64) + threadIdx.x)]*2.000000f)
    }
  }
}

3.6 compute_at

For a schedule that consists of multiple operators, TVM will compute tensors at the root separately by default.[摘自官方文档]

一句话概括compute_at的作用：明明一个for循环能完成的事，为啥要分成多个（注意：是多个而不是多层）for循环串行执行？抓紧合并成一个！

如下的例子中，计算的内容为B = A + 1与C = B * 2，且为element-wise的计算。默认的schedule会使用两个for循环，在第1个for循环中计算B = A + 1，完成后才用新的for循环计算C = B * 2。

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

s = tvm.create_schedule(C.op)
print(tvm.lower(s, [A, B, C], simple_mode=True))

  for (i, 0, m) {
    B[i] = (A[i] + 1.000000f)
  }
}
produce C {
  for (i, 0, m) {
    C[i] = (B[i]*2.000000f)
  }
}

而该计算流程可以在一个for循环中完成，因此这2个for循环需要进行合并，使用s[B]pute_at(s[C], C.op.axis[0])，将这两步操作在一个for循环中实现，实现后的代码如下：

produce C {
  for (i, 0, m) {
    produce B {
      B[i] = (A[i] + 1.000000f)
    }
    C[i] = (B[i]*2.000000f)
  }
}

3.7 compute_inline

compute_inline can mark one stage as inline, then the body of computation will be expanded and inserted at the address where the tensor is required.[摘自官方文档]

一句话概括compute_inline的作用：约减中间步骤，不让中间商赚差价！

下面代码中的例子的作用是B = A + 1与C = B * 2，其运算流程可以约减为C = (A + 1) * 2。所以，相比于3.6 compute_at中的输出结果，for循环中的两步计算，可以使用一步来代替，减少计算量从而提高性能。

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

s = tvm.create_schedule(C.op)
s[B].compute_inline()
print(tvm.lower(s, [A, B, C], simple_mode=True))

使用s[B]pute_inline()后，得到的ir中，for循环只有一步的计算量。

produce C {
  for (i, 0, m) {
    C[i] = ((A[i] + 1.000000f)*2.000000f)
  }
}

3.8 compute_root

compute_root can move computation of one stage to the root.[摘自官方文档]

一句话概括compute_root的作用：让schedule恢复到默认的状态（这个方法过于“神奇”，在开发中还魅有使用过）。

如下面的代码所示，通过使用compute_at，是的计算B的for循环与计算C的for循环合并到一起，但是因为使用了s[B]pute_root()，使得计算B的schedule恢复到原始默认状态，那么得到的结果又变成了2个单独的for循环。。。。。。

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

s = tvm.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))

produce B {
  for (i, 0, m) {
    B[i] = (A[i] + 1.000000f)
  }
}
produce C {
  for (i, 0, m) {
    C[i] = (B[i]*2.000000f)
  }
}

更多推荐

TVM之Schedule