如下代码所示,当定义完计算逻辑后,需要在算子接口实现函数中实现调度与编译。
通过调用auto_schedule接口,便可以自动生成相应的调度。配置信息包括算子内核名以及输入、输出张量。
with tvm.target.cce(): schedule = dsl.auto_schedule(result) config = { "name": kernel_name, "tensor_list": [input_data, result] "bool_storage_as_1bit":True } dsl.build(schedule, config)
schedule可以理解为:描述的计算过程如何在硬件上高效执行。就是把相关的计算和硬件设备上的相关指令对应起来。
注意:输入tensor需要是placeholder接口返回的tensor对象,此tensor对象的内存地址不能被覆盖。
例如:"tensor_list": [tensor_a, tensor_b, res],tensor_a与tensor_b是输入张量,res为输出张量。
当tbe.dsl.vcmp与tbe.dsl.vsel接口的mode/condition为bool时,需要设置此参数为False。
编译完成后,会生成算子目标文件.o文件与算子描述文件.json文件。