编译接口
auto_schedule
auto_schedule接口的作用是把定义好的计算过程生成schedule对象。
您可以在ATC安装目录下的“python/site-packages/topi/generic/cce.py”查看接口定义。
函数原型
topi.generic.auto_schedule(outs, option=None)
参数说明
outs:算子输出tensor的列表,支持单个输出或者多输出的场景。
option:用于RL算子搜索的配置参数,正常auto_schedule无需配置。
返回值:
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)
cce_build_code
对schedule对象进行编译生成算子二进制文件与算子描述文件。
您可以在ATC安装目录下的“python/site-packages/te/lang/cce/te_schedule/cce_schedule.py”查看接口定义。
函数原型
te.lang.cce.cce_build_code(sch, config_map=None)
参数说明
- sch:tvm.schedule,schedule to build or to print lower code。
- config_map:build的参数配置,是一个字典,默认是None,直接使用默认配置,包含如下key。
- print_ir:是否打印lower IR code,默认是True。
- need_build:是否进行build,默认是True。
- name:算子的名字,默认是`cce_op`。
只能是大小写字母、数字、“_”的组合,且必须是字母或者“_”开头,长度小于或等于200个字符。
- tensor_list:算子的输入和输出tensor列表,输入是placeholder接口返回的tensor对象,输出是经过计算后的tensor对象,必填值,否则会报错。而且这个列表决定了生成算子的kernel函数的参数的顺序,和此list中的输入和输出的顺序是一致的。
- bool_storage_as_1bit:Bool类型存储时是否按照1 bit存储。
- True:按照1bit存储。
- False:按照8bit存储。
默认值是True。
一个算子只允许选择一种存储模式。
返回值
无。
调用示例
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)