build接口
build接口的作用是把定义好的计算过程生成schedule对象,并进行build生成cce算子文件。
topi.generic.auto_schedule(outs)
生成dsl的schedule。该接口在cce.py中定义。
参数说明
outs:对算子的计算图描述,就是DSL。
返回值:
schedule:算子的计算schedule。
调用示例
import te.lang.cce from te import tvm import topi.generic shape = (28,28) dtype = "float16" # 定义输入 data = tvm.placeholder(shape, name="data", dtype=dtype) # 描述算子计算过程 res = te.lang.cce.vabs(data) with tvm.target.cce(): # 生成schedule对象 sch = topi.generic.auto_schedule(res)
te.lang.cce.cce_build_code(sch, config_map = {})
对schedule打印lower code或者进行build。该接口在cce_schedule.py中定义。
参数说明
- sch:tvm.schedule,schedule to build or to print lower code。
- config_map:build的参数配置,是一个字典,默认是{}并使用默认配置,包含如下key。
- print_ir:是否打印lower IR code,默认是True。
- need_build:是否进行build,默认是True。
- name:算子的名字,默认是`cce_op`。
- tensor_list:算子的输入和输出tensor列表,输入是placeholder接口返回的tensor对象,输出是经过计算后的tensor对象,必填值,否则会报错。而且这个列表决定了生成算子的kernel函数的参数的顺序,和此list中的输入和输出的顺序是一致的。
返回值
无。
调用示例
import te.lang.cce from te import tvm from topi import generic # 定义输入占位符 data = tvm.placeholder(shape, name="data", dtype=dtype) with tvm.target.cce(): # 描述算子计算过程 res = te.lang.cce.vabs(data) # 生成schedule对象 sch = generic.auto_schedule(res) # 定义build配置参数 config = {"print_ir" : True, "need_build" : True, "name" : "abs_28_28_float16", "tensor_list" : [data,res] } # build算子 te.lang.cce.cce_build_code(sch, config)