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

最新推荐文章于 2023-02-15 11:24:40 发布

BesiseB

最新推荐文章于 2023-02-15 11:24:40 发布

阅读量1.8k

点赞数 1

本文链接：https://blog.csdn.net/B10030607/article/details/115954499

版权

目的

编译器是下一代大厂之争的关键，曾经有过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(op):
    """CUDA lowering rule for log"""
    if op.dtype == "float32":
        return tvm.tir.call_pure_extern("float32", "logf", op.args[0])
    elif op.dtype == "float64":
        return tvm.tir.call_pure_extern("float64", "log", op.args[0])
    else:
        return op
        
# new op registration is triggered by registering an attribute of the op
tvm.ir.register_op_attr("tir.mylog", "TCallEffectKind", tvm.tir.CallEffectKind.Pure)
#tvm.ir.register_op_attr("xxx.mylog", "TCallEffectKind", tvm.tir.CallEffectKind.Pure)
#tvm.ir.register_op_attr("tir.yyy", "TCallEffectKind", tvm.tir.CallEffectKind.Pure)
tvm.target.register_intrin_rule("cuda", "mylog", my_cuda_mylog_rule, override=True)

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

3. 扫描与递归核

递归计算在神经网络中非常典型，TVM如何支持递归计算呢？

TVM支持使用扫描运算符（scan operator）来描述符号循环。

tensor = tvm.te.scan(init, update, state_placeholder, inputs=None, name=‘scan’, tag=’’, attrs=None)
通过沿轴扫描构建新tensor.

3.1 单阶段扫描

下面以沿行计算矩阵的累加和（np.cumsum(a, axis=0)为例:

m = te.var("m")
n = te.var("n")
X = te.placeholder((m, n), name="X")
s_state = te.placeholder((m, n))
s_init = te.compute((1, n), lambda _, i: X[0, i]) // 以X的第一行作为s_init
s_update = te.compute((m, n), lambda t, i: s_state[t - 1, i] + X[t, i]) // s_state的t行等于s_state的t-1行加X的t行
s_scan = tvm.te.scan(s_init, s_update, s_state, inputs=[X]) // 沿t轴扫描X，递归计算得到s_scan
print(tvm.lower(s, [X, s_scan], simple_mode=True))

默认调度下，s_scan的计算过程如下：

for (i: int32, 0, n) {
    scan_2[(i*stride_1)] = (float32*)X_2[(i*stride_3)]
  }
  for (scan.idx: int32, 0, (m - 1)) {
    for (i_1: int32, 0, n) {
      scan_2[(((scan.idx + 1)*stride) + (i_1*stride_1))] = ((float32*)scan_2[((scan.idx*stride) + (i_1*stride_1))] + (float32*)X_2[(((scan.idx + 1)*stride_2) + (i_1*stride_3))])
    }
  }

可以沿列split轴做进一步优化，但是注意，不能沿递归轴切割。

s = te.create_schedule(s_scan.op)
num_thread = 256
block_x = te.thread_axis("blockIdx.x")
thread_x = te.thread_axis("threadIdx.x")
xo, xi = s[s_init].split(s_init.op.axis[1], factor=num_thread)
s[s_init].bind(xo, block_x)
s[s_init].bind(xi, thread_x)
print(tvm.lower(s, [X, s_scan], simple_mode=True))

到这里是对s_init的初始化计算做优化：

attr [IterVar(blockIdx.x: int32, (nullptr), "ThreadIndex", "blockIdx.x")] "thread_extent" = floordiv((n + 255), 256);
  attr [IterVar(threadIdx.x: int32, (nullptr), "ThreadIndex", "threadIdx.x")] "thread_extent" = 256;
  if @tir.likely((((blockIdx.x*256) + threadIdx.x) < n), dtype=bool) {
    scan_2[(((blockIdx.x*256) + threadIdx.x)*stride_1)] = (float32*)X_2[(((blockIdx.x*256) + threadIdx.x)*stride_3)]
  }
  for (scan.idx: int32, 0, (m - 1)) {
    for (i: int32, 0, n) {
      scan_2[(((scan.idx + 1)*stride) + (i*stride_1))] = ((float32*)scan_2[((scan.idx*stride) + (i*stride_1))] + (float32*)X_2[(((scan.idx + 1)*stride_2) + (i*stride_3))])
    }
  }

还可以继续对s_update的递归过程做优化：

xo, xi = s[s_update].split(s_update.op.axis[1], factor=num_thread)
s[s_update].bind(xo, block_x)
s[s_update].bind(xi, thread_x)
print(tvm.lower(s, [X, s_scan], simple_mode=True))

最终的效果如下：

attr [IterVar(blockIdx.x: int32, (nullptr), "ThreadIndex", "blockIdx.x")] "thread_extent" = floordiv((n + 255), 256);
  attr [IterVar(threadIdx.x: int32, (nullptr), "ThreadIndex", "threadIdx.x")] "thread_extent" = 256;
  if @tir.likely((((blockIdx.x*256) + threadIdx.x) < n), dtype=bool) {
    scan_2[(((blockIdx.x*256) + threadIdx.x)*stride_1)] = (float32*)X_2[(((blockIdx.x*256) + threadIdx.x)*stride_3)]
  }
  for (scan.idx: int32, 0, (m - 1)) {
    attr [IterVar(blockIdx.x, (nullptr), "ThreadIndex", "blockIdx.x")] "thread_extent" = floordiv((n + 255), 256);
    attr [IterVar(threadIdx.x, (nullptr), "ThreadIndex", "threadIdx.x")] "thread_extent" = 256;
    if @tir.likely((((blockIdx.x*256) + threadIdx.x) < n), dtype=bool) {
      scan_2[(((scan.idx + 1)*stride) + (((blockIdx.x*256) + threadIdx.x)*stride_1))] = ((float32*)scan_2[((scan.idx*stride) + (((blockIdx.x*256) + threadIdx.x)*stride_1))] + (float32*)X_2[(((scan.idx + 1)*stride_2) + (((blockIdx.x*256) + threadIdx.x)*stride_3))])
    }
  }

3.2 多阶段扫描

TVM也支持多阶段扫描，并且设计了一组约束，以支持compute_at不出错。
举例来说：

m = te.var("m")
n = te.var("n")
X = te.placeholder((m, n), name="X")
s_state = te.placeholder((m, n))
s_init = te.compute((1, n), lambda _, i: X[0, i])
s_update_s1 = te.compute((m, n), lambda t, i: s_state[t - 1, i] * 2, name="s1")
s_update_s2 = te.compute((m, n), lambda t, i: s_update_s1[t, i] + X[t, i], name="s2")
s_scan = tvm.te.scan(s_init, s_update_s2, s_state, inputs=[X])
s = te.create_schedule(s_scan.op)
xo, xi = s[s_update_s2].split(s_update_s2.op.axis[1], factor=32)
print(tvm.lower(s, [X, s_scan], simple_mode=True))

得到：

for (i: int32, 0, n) {
      scan_2[(i*stride_1)] = (float32*)X_2[(i*stride_3)]
    }
    for (scan.idx: int32, 0, (m - 1)) {
      for (i_1: int32, 0, n) {
        s1[i_1] = ((float32*)scan_2[((scan.idx*stride) + (i_1*stride_1))]*2f32)
      }
      for (i.outer: int32, 0, floordiv((n + 31), 32)) {
        for (i.inner: int32, 0, 32) {
          if @tir.likely((((i.outer*32) + i.inner) < n), dtype=bool) {
            scan_2[(((scan.idx + 1)*stride) + (((i.outer*32) + i.inner)*stride_1))] = ((float32*)s1[((i.outer*32) + i.inner)] + (float32*)X_2[(((scan.idx + 1)*stride_2) + (((i.outer*32) + i.inner)*stride_3))])
          }
        }
      }
    }

如果使用compute_at，将第一个stage的计算向第二个stage移动：

s[s_update_s1].compute_at(s[s_update_s2], xo)
print(tvm.lower(s, [X, s_scan], simple_mode=True))

则得到：

for (i: int32, 0, n) {
      scan_2[(i*stride_1)] = (float32*)X_2[(i*stride_3)]
    }
    for (scan.idx: int32, 0, (m - 1)) {
      for (i.outer: int32, 0, floordiv((n + 31), 32)) {
        for (i_1: int32, 0, 32) {
          if @tir.likely((((i.outer*32) + i_1) < n), dtype=bool) {
            s1[i_1] = ((float32*)scan_2[((scan.idx*stride) + (((i.outer*32) + i_1)*stride_1))]*2f32)
          }
        }
        for (i.inner: int32, 0, 32) {
          if @tir.likely((((i.outer*32) + i.inner) < n), dtype=bool) {
            scan_2[(((scan.idx + 1)*stride) + (((i.outer*32) + i.inner)*stride_1))] = ((float32*)s1[i.inner] + (float32*)X_2[(((scan.idx + 1)*stride_2) + (((i.outer*32) + i.inner)*stride_3))])
          }
        }
      }
    }

3.3 多状态扫描

向RNN中的复杂应用，需要多个递归状态，这一点TVM也支持。

4. 外部张量函数

虽然TVM支持透明代码生成，但是有时候也需要在pipeline中嵌入手写代码。
例如，用cuDNN做卷积核，再定义其它stage。

TVM支持以黑匣子的形式调用函数。具体地说，TVM支持所有与DLPack兼容的张量函数。这意味着我们可以使用POD类型（指针pointer，整数int，浮点数float）或指向DLTensor的指针作为参数来调用任何函数。

cblas是矩阵运算库。

4.1 使用外部张量函数

tensor = tvm.te.extern(shape, inputs, fcompute, name=‘extern’, dtype=None, in_buffers=None, out_buffers=None, tag=’’, attrs=None)
通过外部函数计算tensor。
shape–输出的shape, dtype–输出的dtype

例如，下面先调用cblas进行matmul，再用tvm计算加法，组合使用。

n = 1024
l = 128
m = 235
bias = te.var("bias", dtype="float32")
A = te.placeholder((n, l), name="A")
B = te.placeholder((l, m), name="B")
C = te.extern((n, m), [A, B],
    lambda ins, outs: tvm.tir.call_packed(
        "tvm.contrib.cblas.matmul", ins[0], ins[1], outs[0], False, False
    ),
    name="C",
)
D = te.compute(C.shape, lambda i, j: C[i, j] + bias, name="D")
s = te.create_schedule(D.op)

4.2 贡献库的外部封装

Extern Contrib Wrappers
TVM也支持对外部调用进行封装，这部分代码属于tvm.contrib，也就是不在主库中，而是世界各地的开发者贡献的辅助库。
例如，上段代码的cblas.matmul的调用，等效于：

from tvm.contrib import cblas

C = cblas.matmul(A, B)
D = te.compute(C.shape, lambda i, j: C[i, j] + bias, name="D")
s = te.create_schedule(D.op)

4.3 通过钩子调用外部python函数

TVM支持调用任何PackedFunc，所以也支持用外部函数回调到python中。这使得TVM更加灵活。

@tvm.register_func("tvm.contrib.my_tvm_addone")
def my_tvm_addone(x, y):
    print("my_tvm_addone signatures: %s, %s" % (type(x), type(y)))
    tvm.nd.array(x.asnumpy() + 1).copyto(y)

上述代码在python中通过tvm.register_func装饰器，将python函数my_tvm_addone()注册到TVM运行时系统。这样就可以方便的用te.extern调用了：

A = te.placeholder((n,), name="A")
B = te.extern(
    A.shape,
    [A],
    lambda ins, outs: tvm.tir.call_packed("tvm.contrib.my_tvm_addone", ins[0], outs[0]),
    name="C",
)
s = te.create_schedule(B.op)

5. 张量化

通过使用调度原语tensorize，人们可以用相应的内置函数代替计算单元，从而轻松利用手工制作的微内核，并扩展TVM以支持新的硬件体系结构。

5.1 矩阵乘

以计算 $A*B^T$ 为例：

N, M, L = 1024, 512, 64
A = te.placeholder((N, L), name="A")
B = te.placeholder((M, L), name="B")
k = te.reduce_axis((0, L), name="k")
C = te.compute((N, M), lambda i, j: te.sum(A[i, k] * B[j, k], axis=k), name="C")
s = te.create_schedule(C.op)

这就是中规中矩的三层循环，假设现在有一个支持矩阵向量乘的硬件原语，如GEMV，它要求reduce轴以外的轴的维度不大于16。
为了利用这个硬件原语做加速，那么就要重新设计上述循环：

factor = 16
x, y = C.op.axis
(z,) = C.op.reduce_axis
yo, yi = s[C].split(y, factor=factor)
s[C].reorder(x, yo, yi, z)

这样，原来的三层循环i:1024–j:512–k:64，就变成了i:1024–j_outer:512/16–j_inner:16–k:64，那么里面的两层就符合GEMV的要求了。
而且在i固定的情况下，对矩阵A的访问是沿列连续的。

下面先为GEMV定义内置函数，再沿j_inner做张量化。

5.2 定义张量化内置函数

分为两步：

GEMV的计算定义
TVM据此在原生Matmul调度中匹配计算模式；
指定如何在device上执行GEMV
见下面的intrin_func

def intrin_gemv(m, l):
    a = te.placeholder((l,), name="a")
    b = te.placeholder((m, l), name="b")
    k = te.reduce_axis((0, l), name="k")
    c = te.compute((m,), lambda i: te.sum(a[k] * b[i, k], axis=k), name="c")
    Ab = tvm.tir.decl_buffer(a.shape, a.dtype, name="A", offset_factor=1, strides=[1]) # 张量化需要指定offset_factor，指定计算数据和原始数据起始地址的偏移
    Bb = tvm.tir.decl_buffer(b.shape, b.dtype, name="B", offset_factor=1, strides=[te.var("s1"), 1]) # B是2D，所以指定的步长stride是两个数
    Cb = tvm.tir.decl_buffer(c.shape, c.dtype, name="C", offset_factor=1, strides=[1])

    def intrin_func(ins, outs):
        ib = tvm.tir.ir_builder.create() # 这个方法没找到doc
        aa, bb = ins
        cc = outs[0]
        ib.emit(
            tvm.tir.call_extern(
                "int32", # 输出的数据类型
                "gemv_update",  # 外部函数名
                cc.access_ptr("w"), # 接受输入和输出，将它们转换为指针，并发出（emit）外部函数调用
                aa.access_ptr("r"),
                bb.access_ptr("r"),
                m,
                l,
                bb.strides[0],
            )
        )
        return ib.get()

    return te.decl_tensor_intrin(c.op, intrin_func, binds={a: Ab, b: Bb, c: Cb})

Buffer = tvm.tir.decl_buffer(shape, dtype=None, name=‘buffer’, data=None, strides=None, elem_offset=None, scope=’’, data_alignment=- 1, offset_factor=0, buffer_type=’’, span=None)
声明符号化缓冲区（buffer）。
一般而言，在调用tvm.lower或者tvm.build的时候会自动创建buffer。
当我们想要指定自己的缓冲区布局（layout）时，才用到这个接口。
那strides=[te.var("s1"), 1]是什么意思呢？TVM知道B是紧凑的，所以可以自己推导出来步长是[L, 1]

TensorIntrin = tvm.te.decl_tensor_intrin(op, fcompute, name=‘tensor_intrin’, binds=None, scalar_params=None, default_buffer_params=None)
声明一个张量内置函数。

5.3 张量化

现在可以用外部函数gemv_update做内两层循环的向量化：

gemv = intrin_gemv(factor, L) # factor=16, L=64
s[C].tensorize(yi, gemv)
print(tvm.lower(s, [A, B, C], simple_mode=True))

得到：

for (i: int32, 0, 1024) {
    for (j.outer: int32, 0, 32) {
      @tir.call_extern("gemv_update", @tir.tvm_access_ptr(@tir.type_annotation(, dtype=float32), C_2, ((i*512) + (j.outer*16)), 16, 2, dtype=handle), @tir.tvm_access_ptr(@tir.type_annotation(, dtype=float32), A_2, (i*64), 64, 1, dtype=handle), @tir.tvm_access_ptr(@tir.type_annotation(, dtype=float32), B_2, (j.outer*1024), 1024, 1, dtype=handle), 16, 64, 64, dtype=int32)
    }
  }

为了能build、run这个module，还需要做其他事，比如用stage的pragma导入等，时间有限，先跳过。

6. 输入元组

使用输入元组（tuple），可以进行多个相同shape输入的计算或者约简。
元组的元素分别定义，然后用小括号组织在一起，形成tuple。

有时，我们需要多个输入来表示某些约简运算符，例如argmax，这个场景下，输入元组的各元素将一起协作。在简化过程中，argmax需要比较操作数的值，还需要保留操作数的索引。

Class tvm.tir.Select(condition, true_value, false_value, span=None)
用于选择。

function = tvm.te.comm_reducer(fcombine, fidentity, name=‘reduce’)
创建一个可交换的约简器以进行约简。
fcombine–二元函数，接收两个输入表达式，返回一个输出表达式；
fidentity–接收一个string，返回一个常量表达式；

时间有限，先跳过。

7. TEDD与可视化张量表达式

张量表达式的调度是基于原语的。虽然单个原语的功能容易理解，但是复杂组合后的调度就很难直观理解了。
除了把调度lower成文本进行理解之外，，还可以通过TEDD工具可视化。

Tensor Expression中引入了调度原语的可操作模型 operational model ：

不同计划原语之间的相互作用
调度原语对最终代码生成的影响

这个可操作模型基于数据流图 Dataflow Graph，调度树 Schedule Tree 和迭代变量关系图 IterVar Relationship Graph。调度原语对这些图执行操作。

TEDD根据给定的调度渲染这三个图形。

首先要安装graphviz工具：
sudo apt-get install graphviz
pip3 install graphviz
python中保存的.dot文件可以用dot工具转换为.png图片。
dot /tmp/dfg.dot -Tpng -o ./dfg.png

batch = 1
in_channel = 256
in_size = 32
num_filter = 1024
kernel = 3
stride = 1
padding = "SAME"
dilation = 1

A = te.placeholder((in_size, in_size, in_channel, batch), name="A")
W = te.placeholder((kernel, kernel, in_channel, num_filter), name="W")
B = te.placeholder((1, num_filter, 1), name="bias")

with tvm.target.Target("llvm"):
    t_conv = topi.nn.conv2d_hwcn(A, W, stride, padding, dilation)
    t_bias = topi.add(t_conv, B)
    t_relu = topi.nn.relu(t_bias)
    s = topi.generic.schedule_conv2d_hwcn([t_relu])

对以上脚本，我们分别产生上述三个图看看：

tedd.viz_dataflow_graph(s, dot_file_path="/tmp/dfg.dot")
tedd.viz_schedule_tree(s, dot_file_path="/tmp/scheduletree.dot")
s = s.normalize()
tedd.viz_schedule_tree(s, dot_file_path="/tmp/scheduletree2.dot")
tedd.viz_itervar_relationship_graph(s, dot_file_path="/tmp/itervar.dot")

在这里插入图片描述
数据流图中每个节点代表一个stage，椭圆表示输入输出数据的shape和type，箭头表示数据依赖方向。

在这里插入图片描述
调度图中，ROOT下的每个块代表一个stage，stage名称显示在顶部，计算显示在底部。中间的用于迭代向量，外部越高，内部越低。迭代变量行包含：索引，名称，类型，和范围等其他信息。

在这里插入图片描述
迭代变量关系图中，每个子图代表一个stage，并包含IterVar节点和转换节点。

8. 约简

本节介绍在TVM中实现约简，典型的约简运算符包括sum/max/min/cumsum等。

B = numpy.sum(A, axis=1)
A nxm, B nx1

对应的TVM代码:

n = te.var("n")
m = te.var("m")
A = te.placeholder((n, m), name="A")
k = te.reduce_axis((0, m), "k")
B = te.compute((n,), lambda i: te.sum(A[i, k], axis=k), name="B")

约简的一个问题是我们不能简单地在约简轴上并行化。我们需要对约简的计算进行划分，在对临时数组进行约简之前，将局部约简结果存储在临时数组中。

rfactor原语对计算进行这种重写。在下面的调度中，将B的结果写入临时结果B.rf。

s = te.create_schedule(B.op)
ko, ki = s[B].split(B.op.reduce_axis[0], factor=16)
BF = s.rfactor(B, ki)
print(tvm.lower(s, [A, B], simple_mode=True))

A的维度是nxm，B的约简轴是m轴/列；
split将计算从 n–m的两层循环，切分成n–m/16–16的三层循环，ki就是16那个轴；
rfactor原语进行临时变量的重写，BF的计算是16–n--m/16的三层循环，B.rf的维度是16xn。
在这里插入图片描述
有了B.rf的16xn，计算B就是将16这个维度做累加。
这样做的好处是，B的列的累加可以拆分成m/16个子矩阵的累加和最后的累加，子矩阵的累加之间没有依赖，可以并行化。

for (k.inner: int32, 0, 16) {
      for (i: int32, 0, n) {
        B.rf[((k.inner*n) + i)] = 0f32
        for (k.outer: int32, 0, floordiv((m + 15), 16)) {
          if @tir.likely((((k.outer*16) + k.inner) < m), dtype=bool) {
            B.rf[((k.inner*n) + i)] = ((float32*)B.rf[((k.inner*n) + i)] + (float32*)A_2[((i*stride_1) + (((k.outer*16) + k.inner)*stride_2))])
          }
        }
      }
    }
    for (ax0: int32, 0, n) {
      B_2[(ax0*stride)] = 0f32
      for (k.inner.v: int32, 0, 16) {
        B_2[(ax0*stride)] = ((float32*)B_2[(ax0*stride)] + (float32*)B.rf[((k.inner.v*n) + ax0)])
      }
    }

class tvm.te.Schedule的成员函数rfactor:
tensor = rfactor(tensor, axis, factor_axis=0)
将调度中的约简轴分解为显式轴。
这将创建一个新stage，该stage生成以轴为第一维的新tensor。

后面还有2D约简实现卷积，时间有限，先跳过。

BesiseB

关注

1
点赞
踩
11

收藏

觉得还不错? 一键收藏
0
评论
【tvm官网教程】张量表达与调度

【tvm官网教程】张量表达与调度目的1. 调度原语1.1 te常用接口1.2 stage常用成员函数欢迎使用Markdown编辑器新的改变功能快捷键合理的创建标题，有助于目录的生成如何改变文本的样式插入链接与图片如何插入一段漂亮的代码片生成一个适合你的列表创建一个表格设定内容居中、居左、居右SmartyPants创建一个自定义列表如何创建一个注脚注释也是必不可少的KaTeX数学公式新的甘特图功能，丰富你的文章UML 图表FLowchart流程图导出与导入导出导入目的编译器是下一代大厂之争的关键，曾经有过
复制链接

扫一扫