【tvm官网教程】张量表达与调度

目的

编译器是下一代大厂之争的关键，曾经有过tvm开发基础，现在针对新教程做系统学习。

【tvm官网教程】设计和架构

【tvm官网教程】tvm入门

【tvm官网教程】张量表达与调度
【tvm官网教程】TOPI:TVM算子清单

【tvm官网教程】优化张量算子
【tvm官网教程】AutoTVM:基于模板的自动调优
【tvm官网教程】AutoSchedule:无模板的自动调度

【tvm官网教程】编译DL模型

【tvm官网教程】开发者教程

tvm.tir – Namespace for Tensor-level IR
tvm.ir – Common data structures across all IR variants.
tvm.target – Target description and codgen module.
tvm.te – Namespace for Tensor Expression Language.
tvm.contrib – Contrib APIs of TVM python package.Contrib API provides many useful not core features.
tvm.topi – TVM Operator Inventory.

1. 调度原语

TVM is a domain specific language for efficient kernel construction.
TVM是一种用于高效构建内核的领域专用语言（DSL）。

一个数学上的计算，在计算机上的实现可以用多种方式，不同实现方式的内存访问、性能也不同。因此，TVM要求用户提供“调度”，来明确指定如何开展计算。

调度（Schedule）是计算的变换的组合，它通过变化程序的计算循环Loop，实现不同的性能。
一个调度由多个阶段（Stage）组成，一个 阶段表示一个操作（operation）的调度。原语提供各种方法来调度每个阶段。

1.1 te常用接口

var = tvm.te.var(name=‘tindex’, dtype=‘int32’, span=None)
创建符号化变量，注意，只是“符号”，不占内存。

tensor = tvm.te.placeholder(shape, dtype=None, name=‘placeholder’)
创建空tensor对象，注意，只是“占位符”，不占内存。

tensor = tvm.te.compute(shape, fcompute, name=‘compute’, tag=’’, attrs=None)
通过lambda表达式，创建一个指定shape的新tensor。

schedule = tvm.te.create_schedule(ops)
创建计算表达式ops(操作列表)的一个调度，默认串行方式、行主序（也就是MN的N在内存中连续）方式计算。

IterVar = tvm.te.thread_axis(dom=None, tag=’’, name=’’, span=None)
创建一个表示线程下标的迭代变量，常与stage原语bind配合使用。

IterVar = tvm.te.reduce_axis(dom, name=‘rv’, thread_tag=’’, span=None)
创建一个用于reduction的迭代变量。reduction是约简，意味着沿某个轴做累加累积等操作，最终维度会降低，如从2D变为1D。

PrimExpr = tvm.te.exp(x)
屏蔽硬件和数据类型，提供统一的指数计算接口。

1.2 tvm常用接口

IRModule = tvm.lower(sch, args, name=‘main’, binds=None, simple_mode=False)
在调度被build到特定target之前，将调度向底层表达成IRModule。

model = tvm.build(inputs, args=None, target=None, target_host=None, name=‘default_function’, binds=None)
构建一个带参数签名的函数，为与目标信息耦合的设备生成代码。

• 当TVM编译设备专用程序（例如CUDA）时，我们还需要主机（CPU）侧代码才能与驱动程序进行交互，以正确设置尺寸和参数。target_host用于指定主机端代码生成目标。默认情况下，如果启用了llvm，则使用它，否则使用stackvm解释器。
• 返回的model中包含host侧代码和device侧代码。

1.3 stage常用成员函数

outer, inner = split(parent, factor=None, nparts=None)
将stage分割成多份，factor指定内循环，nparts指定外循环。例如一层20的循环，变成外层4、内层5的两层循环。
parent的数据类型是IterVar，迭代变量；

例如矩阵乘场景中，B的访问模式不连续，cache命中率低，将B拆成3D则能优化这一点。

x_outer, y_outer, x_inner, y_inner = tile(x_parent, y_parent, x_factor, y_factor)
在两个维度进行分块，可以看做是2D版本的split。

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)
// 上一句等效于下三句
/* xo, xi = s[B].split(B.op.axis[0], factor=10)
yo, yi = s[B].split(B.op.axis[1], factor=5)
s[B].reorder(xo, yo, xi, yi) */

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

fused = fuse(*args) args–List of IterVars
将多个连续的迭代变量合并成一个。

reorder(*args)
以指定顺序重排迭代变量。

bind(ivar, thread_ivar)
将迭代变量ivar绑定到指定线程，这个通常用于GPU编程。

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)
print(tvm.lower(s, [A, B], simple_mode=True))

s[B].bind(bx, te.thread_axis("blockIdx.x"))
s[B].bind(tx, te.thread_axis("threadIdx.x"))
print(tvm.lower(s, [A, B], simple_mode=True))

for (i.outer: int32, 0, floordiv((n + 63), 64)) {
   
    for (i.inner: int32, 0, 64) {
   
      if @tir.likely((((i.outer*64) + i.inner) < n), dtype=bool) {
   
        B_2[(((i.outer*64) + i.inner)*stride)] = ((float32*)A_2[(((i.outer*64) + i.inner)*stride_1)]*2f32)
      }
    }
  }
####### 使用的bind原语之后：
attr [IterVar(blockIdx.x: int32, (nullptr), "ThreadIndex", "blockIdx.x")] "thread_extent" = floordiv((n + 63), 64);
  attr [IterVar(threadIdx.x: int32, (nullptr), "ThreadIndex", "threadIdx.x")] "thread_extent" = 64;
  if @tir.likely((((blockIdx.x*64) + threadIdx.x) < n), dtype=bool) {
   
    B_2[(((blockIdx.x*64) + threadIdx.x)*stride)] = ((float32*)A_2[(((blockIdx.x*64) + threadIdx.x)*stride_1)]*2f32)
  }

compute_at(parent, scope) parent–Stage, scope–The loop scope to be attached to
对于包含多个运算符（operator）的调度，默认情况下，TVM将在根处分别计算张量。
compute_at将该stage的计算移动到parent stage的scope轴，减少一层循环。

compute_inline()
标记该stage为内联形式，在后续需要它的时候，再将具体计算展开并插入。

compute_root()
与compute_at相反，将该stage的计算放到根处。

2. 内置函数与数学函数

有些复杂函数，如exp()，其实现取决于目标系统，而且有可能在不同target平台的名称也不同。
本节学习，如何用TVM统一的API接口调用这些“目标特定函数”。

2.1 直接声明外部数学调用

例如，调用仅在CUDA支持的浮点型数字的指数运算函数__expf：

n = te.var("n")
A = te.placeholder((n,), name="A")
B = te.compute(A.shape, lambda i: tvm.tir.call_pure_extern("float32", "__expf", A[i]), name="B")

PrimExpr = tvm.tir.call_extern(dtype, func_name, *args, span=None)
通过调用外部函数build表达式。

PrimExpr = tvm.tir.call_pure_extern(dtype, func_name, *args, span=None)
通过调用纯外部函数build表达式。

PrimExpr = tvm.tir.call_intrin(dtype, func_name, *args, span=None)
通过调用内置函数build表达式。

PrimExpr = tvm.tir.call_packed(*args, span=None)
通过调用外部packed函数build表达式。常与te.extern()配合使用，通过调用外部函数创建tensor。

2.2 统一内置函数调用

TVM内置函数(intrinsic)提供了一种机制，使得为任何device、任何数据类型生成相同的代码。这是解决这类问题的推荐的方法。

例如，调用TVM内置函数tvm.te.exp()进行指数运算，在cuda下将调用__expf()，在opencl下将调用exp()：

n = te.var("n")
A = te.placeholder((n,), name="A")
B = te.compute(A.shape, lambda i: te.exp(A[i]), name="B")

2.3 内置函数下降规则

当tvm.te.exp()被调用时，TVM创建一个固有的Call Expr。TVM使用转换规则将内置函数调用转换为设备特定的外部调用。
TVM也支持定制运行时(runtime)的下降(lowering)规则。

function = tvm.target.register_intrin_rule(target, intrin, f=None, override=False)
注册一个内置函数的生成规则，其实就是注册一个回调函数。

• intrin，str类型，内置函数名称；
• f，function类型，想注册的生成规则函数；
• 内置函数生成规则是代码生成器获得设备特定调用的回调函数，
  register_func(“tvm.intrin.rule.%s.%s” % (target, intrin), f, override)

TVM在后端当然已经注册好了许多内置函数生成规则/回调函数，同时TVM也支持用户用这个接口新增、修改规则。

新增规则前，要先用register_op_attr接口注册操作符属性，以触发新操作符的注册。
function = tvm.ir.register_op_attr(op_name, attr_key, value=None, level=10)
通过名称注册运算符（operator）的运算符属性。
根据下面的测试可知:

• register_op_attr的op_name必须以“tir.”打头
• register_intrin_rule的intrin必须有对应的不带“tir."头部的op_name

def balabala(x):
    """customized log intrinsic function"""
    return tvm.tir.call_intrin(x.dtype, "tir.mylog", x)
    #return tvm.tir.call_intrin(x.dtype, "xxx.mylog", x) #RecursionError: maximum recursion depth exceeded
    #return tvm.tir.call_intrin(x.dtype, "tir.yyy", x) #TVMError: Unresolved call Op(tir.yyy)

def my_cuda_mylog_rule