【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(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 A*B^T ABT为例:

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约简实现卷积,时间有限,先跳过

  • 0
    点赞
  • 0
    评论
  • 4
    收藏
  • 一键三连
    一键三连
  • 扫一扫,分享海报

©️2021 CSDN 皮肤主题: 大白 设计师:CSDN官方博客 返回首页
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、C币套餐、付费专栏及课程。

余额充值