TVM Compiler中文教程：TVM调度原语（Schedule Primitives）

admin管理员组

文章数量:1794759

TVM Compiler中文教程：TVM调度原语（Schedule Primitives）

文章目录

• TVM调度原语（Schedule Primitives）
• • 分裂split
  • 平铺tile
  • 融合fuse
  • 重排序reorder
  • 绑定bind
  • 从哪里开始计算compute_at
  • 计算内联compute_inline
  • compute_root
  • 总结

TVM调度原语（Schedule Primitives）

TVM是用于高效内核代码构建的版本领域专用语言（Domain-Specialed-Language，DSL） 。

这篇教程，我们将展示通过TVM提供的各种原语怎么去调度计算。

from __future__ import absolute_import, print_function import tvm import numpy as np

通常存在几种计算相同结果的方法，但是，不同的方法将导致不同的局部性和性能，所以TVM要求用户提供怎么去调用Schedule描述计算是如何执行的。

Schedule是计算的变换的集合，它通过变换程序中的计算循环Loop，实现不同性能。

#定义一些变量 n = tvm.var('n') m = tvm.var('m')

调度（Schedule）能通过一系列计算ops来定义，默认情况下，调度计算张量以航为顺序。

#定义矩阵元素element-wise乘法 A = tvm.placeholder((m,n), name='A') B = tvm.placeholder((m,n), name='B') C = tvmpute((m,n),lambda i,j: A[i,j] * B[i,j], name ='C') #创建调度 s = tvm.create_schedule([C.op]) #lower会将计算从定义转换为真正的可调用函数。 使用参数`simple_mode = True`，它将返回一个可读的C伪代码，我们在这里使用它来打印计划结果。 print(tvm.lower(s, [A, B, C], simple_mode=True))

输出：

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

一个调度过程由多个阶段组成，一个阶段表示操作的一个调度。我们提供各种方法来调度每个阶段。

分裂split

split通过factor分裂指定轴为两个轴。

A = tvm.placeholder((m,), name='A') B = tvmpute((m,), lambda i: A[i]*2, name='B') s = tvm.create_schedule(B.op) #分裂0轴为两个轴，先计算内循环再计算外循环，xo为外循环，xi为内循环 xo, xi = s[B].split(B.op.axis[0], factor=32) print(tvm.lower(s, [A, B], simple_mode=True))

输出：

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) } } } }

使用nparts与factor作用相反，nparts指定外循环次数，factor指定内循环次数。

A = tvm.placeholder((m,), name='A') B = tvmpute((m,), lambda i: A[i], name='B') s = tvm.create_schedule(B.op) bx, tx = s[B].split(B.op.axis[0], nparts=32) print(tvm.lower(s, [A, B], simple_mode=True))

输出：

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

tile通过平铺两个轴执行计算图块

A = tvm.placeholder((m, n), name='A') B = tvmpute((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))

输出：

produce B { for (i.outer, 0, ((m + 9)/10)) { for (j.outer, 0, ((n + 4)/5)) { //先执行10x5的图块，滑动下一个图块 for (i.inner, 0, 10) { for (j.inner, 0, 5) { if (likely(((i.outer*10) < (m - i.inner)))) { if (likely(((j.outer*5) < (n - j.inner)))) { B[(((j.outer*5) + (((i.outer*10) + i.inner)*n)) + j.inner)] = A[(((j.outer*5) + (((i.outer*10) + i.inner)*n)) + j.inner)] } } } } } } } 融合fuse

fuse能融合一个计算的两个轴

A = tvm.placeholder((m,n),name='A') B = tvmpute((m,n), lambda i,j: A[i,j], name='B') s = tvm.create_schedule(B.op) #首先平铺成4轴（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) #然后融合(i.inner,j.inner)进一个轴：（i.inner.j.inner.fused） fused = s[B].fuse(xi,yj) 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) < (m - (i.inner.j.inner.fused/5))))) { if (likely(((j.outer*5) < (n - (i.inner.j.inner.fused % 5))))) { B[(((j.outer*5) + (i.inner.j.inner.fused % 5)) + (((i.outer*10) + (i.inner.j.inner.fused/5))*n))] = A[(((j.outer*5) + (i.inner.j.inner.fused % 5)) + (((i.outer*10) + (i.inner.j.inner.fused/5))*n))] } } } } } } 重排序reorder

reorder能按照指定顺序重新排列轴（类似于permute）。

A = tvm.placeholder((m, n), name='A') B = tvmpute((m, n), lambda i, j: A[i, j], name='B') s = tvm.create_schedule(B.op) #首先平铺成4轴（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) s[B].reorder(xi,yo,xo,yi) print(tvm.lower(s, [A, B], simple_mode=True))

输出：

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) < (m - i.inner)))) { if (likely(((j.outer*5) < (n - j.inner)))) { B[(((j.outer*5) + (((i.outer*10) + i.inner)*n)) + j.inner)] = A[(((j.outer*5) + (((i.outer*10) + i.inner)*n)) + j.inner)] } } } } } } } 绑定bind

bind可以使用线程轴绑定指定的轴，通常在GPU编程中使用。

A = tvm.placeholder((n,), name='A') B = tvmpute(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))

输出：

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) < (n - threadIdx.x)))) { B[((blockIdx.x*64) + threadIdx.x)] = (A[((blockIdx.x*64) + threadIdx.x)]*2.000000f) } } 从哪里开始计算compute_at

对于包含多个算子的调度，TVM默认从root开始遍历计算张量。

A = tvm.placeholder((m,), name='A') B = tvmpute((m,), lambda i: A[i]+1, name='B') C = tvmpute((m,), lambda i: B[i]*2, name='C') s = tvm.create_schedule(C.op) 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) } }

compute_at可以将B的计算移动到C的第一个计算轴。

A = tvm.placeholder((m,), name='A') B = tvmpute((m,), lambda i: A[i]+1, name='B') C = tvmpute((m,), lambda i: B[i]*2, name='C') s = tvm.create_schedule(C.op) # 移动B循环到C循环的第一个轴 s[B]pute_at(S[C], C.op.axis[0]) print(tvm.lower(s, [A, B, C], simple_mode=True))

输出：

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

compute_inline可以将一个计算阶段标记为内联，然后将计算体扩展并插入需要张量的地址处。（和C中的内联函数一个意思）

A = tvm.placeholder((m,), name='A') B = tvmpute((m,), lambda i: A[i]+1, name='B') C = tvmpute((m,), lambda i: B[i]*2, name='C') s = tvm.create_schedule(C.op) s[B]pute_inline() print(tvm.lower(s, [A, B, C], simple_mode=True))

输出：

produce C { for (i, 0, m) { //类似内联函数，直接合成一个循环 C[i] = ((A[i]*2.000000f) + 2.000000f) } } compute_root

compute_root可以将一个计算阶段的计算移动到root。（compute_at的逆过程）

A = tvm.placeholder((m,), name='A') B = tvmpute((m,), lambda i: A[i]+1, name='B') C = tvmpute((m,), lambda i: B[i]*2, name='C') s = tvm.create_schedule(C.op) # B移动到C的0轴 s[B]pute_at(s[C], C.op.axis[0]) # B重新移动回root s[B]pute_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中的调度原语，它允许用户轻松灵活地调度计算。

为了获得良好性能的内核实现，一般工作流程通常是：

• 通过一系列操作描述您的计算。
• 尝试使用原语安排计算。
• 编译并运行以查看性能差异。
• 根据运行结果调整你的调度。

本文标签： 中文教程CompilerTVMschedule