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

调度与编译

如下代码所示,当定义完计算逻辑后,需要在算子接口实现函数中实现调度与编译。

通过调用auto_schedule接口,便可以自动生成相应的调度,此处通过TVM的打印机制可以看到相应计算的中间表示。配置信息包括是否需要打印IR、是否编译以及算子内核名以及输入、输出张量。

with tvm.target.cce():
    schedule = generic.auto_schedule(result)
config = {
    "print_ir": True,
    "need_build": True,
    "name": kernel_name,
    "tensor_list": [input_data, result]
    "bool_storage_as_bit":True
}
te.lang.cce.cce_build_code(schedule, config)
  • 使用“generic”“auto_schedule”接口,自动生成相应的调度(schedule),“auto_schedule”接口的参数为算子的输出张量。

    schedule可以理解为:描述的计算过程如何在硬件上高效执行。就是把相关的计算和硬件设备上的相关指令对应起来。schedule对象中包含一个“中间表示”(IR),它用一种类似伪代码来描述计算过程,可以通过“print_ir”参数把它打印出来进行查看。

  • “need_build”:是否进行build,默认是True。
  • “name”:编译生成的算子二进制文件名称,只能是大小写字母、数字、“_”的组合,且必须是字母或者“_”开头,长度小于或等于200个字符。
  • “tensor_list”(张量列表)中保存输入张量、输出张量,这个顺序需要严格按照算子本身的输入、输出数据顺序排列。

    注意:输入tensor需要是placeholder接口返回的tensor对象,此tensor对象的内存地址不能被覆盖。

    例如:"tensor_list": [tensor_a, tensor_b, res],tensor_a与tensor_b是输入张量,res为输出张量。

  • “bool_storage_as_1bit”:Bool类型存储时是否按照1 bit存储。True:按照1bit存储,False:按照8bit进行存储,默认值为True。

    当te.lang.cce.vcmp(lhs, rhs, op, mode)接口的mode为bool时,需要设置此参数为False。

  • 根据调度和配置使用“te.lang.cce”提供的“cce_build_code”接口来进行算子编译,算子编译过程会根据输入的数据形状、类别、算子参数等编译出专用内核,这个过程在离线模型生成器转换模型时发生。
    • schedule:生成的算子计算schedule对象。
    • config:编译参数配置的map。

    编译完成后,会生成算子目标文件.o文件与算子描述文件.json文件。

分享:

    相关文档

    相关产品

close