更新时间:2021-03-18 GMT+08:00
分享

Auto Schedule

基于TBE DSL编写一个算子,就是通过组合DSL接口表达算子的计算逻辑,然后调用Auto Schedule进行算子的自动调度,完成数据切块和数据流向的划分。Auto Schedule机制是TBE底层的默认Schedule调优机制,开发者无法在算子开发代码过程中进行控制,下面简要介绍Auto Schedule的原理。

如下是基于DSL进行算子开发的示例,实现对x取指数,然后在轴0上进行累加降维,再取倒数的功能。

x = tvm.placeholder((512, 1024), "float16")
exp_x = te.lang.cce.vexp(x)
reduce_exp_x = te.lang.cce.sum(exp_x, axis = 0)
res = te.lang.cce.vrec(reduce_exp_x)

with tvm.target.cce():
    sch = topi.generic.auto_schedule(res)

开发者调用topi.generic.auto_schedule()接口开启TBE的自动调度,自动调度的总体流程如图1所示。

图1 Auto Schedule总体流程
  1. 调用Auto Schedule接口的时候,就是传递了一个compute的语法树,TBE中每一个compute语句在进行编译的时候都会被加上tag_scope标志,如下所示。

    添加tag_scope的操作如下:

    with tvm.tag_scope(op):
        tmp = tvm.compute(shape, lambda_func, name=name)
    图2所示,左侧的compute语法树,也叫抽象语法树(Abstract Syntax Tree: AST),编译过程中,会对每一个compute语句加上tag_scope标志。
    图2 计算语句与tag_scope标志对应示例
  2. 根据scope标识,识别出对应pattern,TBE当前支持的pattern类型有:elewise、reduce、segment、concat、conv、depthwise、pooling2d等。TBE会按照pattern规则对AST进行切分,例如最简单的一条pattern规则是elewise可以和其他的pattern连在一起,reduce、segment、concat不能在一个AST子图内。
  3. 完成AST子图切分后,TBE会创建并初始化Schedule对象。
  4. Schedule执行过程中首先找到AST子图的边界,然后对每一个子图根据其pattern选择一个合适的Schedule模板进行调度。调度过程主要包括数据流管理、tiling以及指令映射等。
分享:

    相关文档

    相关产品

close