为了方便开发者进行自定义算子开发,TBE(Tensor Boost Engine)提供了一套计算接口供开发者用于组装算子的计算逻辑,这套计算接口称之为DSL(Domain-Specific Language)。基于DSL开发的算子,可以直接使用TBE提供的Auto Schedule机制,自动完成调度过程,省去最复杂的调度编写过程。
TBE DSL算子的功能框架如图1所示。
代码示例如下所示:
//初始化输入tensor,为输入tensor进行占位 data_x = tvm.placeholder(shape_x, name="data_1", dtype=input_data_type) data_y = tvm.placeholder(shape_y, name="data_2", dtype=input_data_type) //调用计算接口实现data_x + data_y res = tbe.dsl.vadd(data_x, data_y) //调用auto_schedule接口实现自动调度 with tvm.target.cce(): schedule = tbe.dsl.auto_schedule(res) //配置编译参数并进行编译 config = {"name": kernel_name, "tensor_list": (data_x, data_y, res)} tbe.dsl.build(schedule, config)
Schedule(调度)就是为了解决如上问题,通过调整计算逻辑、优化计算过程,使得计算更加高效,并保证计算过程中占用的硬件存储空间不会超过上限。
TBE DSL提供了Auto Schedule机制,以上Schedule操作无需用户关注。用户通过组合DSL接口表达算子的计算逻辑后,直接调用Auto Schedule接口即可实现自动调度,完成数据切块和数据流向的划分。Auto Schedule机制是TBE底层的默认Schedule调优机制,开发者无法在算子开发代码过程中进行控制,下面简要介绍Auto Schedule的原理。
如下是基于DSL进行算子开发的示例,实现对x取指数,然后在轴0上进行累加降维,再取倒数的功能。
x = tvm.placeholder((512, 1024), "float16") exp_x = tbe.dsl.vexp(x) reduce_exp_x = tbe.dsl.sum(exp_x, axis = 0) res = tbe.dsl.vrec(reduce_exp_x) with tvm.target.cce(): sch = tbe.dsl.auto_schedule(res)
开发者调用tbe.dsl.auto_schedule接口开启TBE的自动调度,自动调度的总体流程如图2所示。
with tvm.tag_scope(op): tmp = tvm.compute(shape, lambda_func, name=name)