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

算子调度与编译

如下代码所示,当定义完计算逻辑后,使用auto_schedule机制,便可以自动生成相应的调度,此处通过TVM的打印机制可以看到相应计算的中间表示。配置信息包括是否需要打印、编译以及算子内核名以及输入、输出张量。

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

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

  • “tensor_list”(张量列表)中保存输入张量、输出张量,这个顺序需要严格按照算子本身的输入、输出数据顺序排列。

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

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

    编译完成后,会生成算子目标文件.o文件(运行目标为AI Core的算子)或者.so文件(运行目标为AI CPU的算子)与算子描述文件.json文件。

分享:

    相关文档

    相关产品