昇腾社区首页
中文
注册

自定义算子入图

当开发者在执行PyTorch/TensorFlow网络或者基于算子原型构造Ascend Graph时,如果遇到了“算子加速库”中不支持的算子,需要开发者先开发自定义算子,再将自定义算子入图。

本章节给出这几种场景下自定义算子开发及入图的整体流程。

PyTorch网络自定义算子开发

完成自定义Ascend C算子开发并进行PyTorch适配和入图执行的具体流程如下图所示:

  1. 自定义算子开发。
    自定义算子入图前,首先需要基于Ascend C完成自定义算子的开发,包括如下几个步骤。
    1. 创建算子工程。
    2. 实现算子,进行算子的原型定义、Kernel侧算子实现与Host侧Tiling实现。
    3. 算子入图开发,主要进行shape推导等算子入图适配函数的实现。
    4. 算子工程编译部署,编译生成自定义算子安装包并进行算子包的安装,将自定义算子部署到算子加速库。

    上述流程的详细描述可参见Ascend C算子开发指南中的算子入图(GE)图开发章节。

  2. PyTorch侧适配。
    自定义算子开发完成后,需要在PyTorch侧进行适配,以实现自定义算子入图,包括如下操作:
    1. 通过OpPlugin完成算子在PyTorch框架注册。
    2. 在torch_npu中注册meta函数,以支持入图时的shape与dtype推导。
    3. 通过TorchAir完成自定义算子converter注册和实现。

    详细的自定义算子入PyTorch图方法请参见PyTorch图模式使用指南(TorchAir)的“自定义算子插件化入图”章节。

TensorFlow网络自定义算子开发

完成自定义Ascend C算子开发并进行TensorFlow框架适配和入图执行的具体流程如下图所示:

  1. 自定义算子开发。
    自定义算子入图前,首先需要基于Ascend C完成自定义算子的开发,包括如下几个步骤。
    1. 创建算子工程。
    2. 实现算子,进行算子的原型定义、Kernel侧算子实现与Host侧Tiling实现。
    3. 算子入图开发,主要进行shape推导等算子入图适配函数的实现。

    上述流程的详细描述可参见Ascend C算子开发指南中的算子入图(GE)图开发章节。

  2. TensorFlow框架适配插件开发。通过REGISTER_CUSTOM_OP注册自定义算子并完成TensorFlow算子到CANN算子的映射。对于TensorFlow自定义算子映射到CANN算子的场景,还需要完成TensorFlow自定义算子的开发,详细说明参见Ascend C算子开发指南中的AI框架算子适配 > TensorFlow框架章节。
  3. 算子工程编译部署,编译生成自定义算子安装包并进行算子包的安装,将自定义算子部署到算子加速库。
  4. TensorFlow框架算子调用。

Ascend Graph构图

自定义Ascend C算子并将其入Ascend Graph执行的开发流程如下图所示:

  1. 自定义算子开发。
    自定义算子入图前,首先需要基于Ascend C完成自定义算子的开发,包括如下几个步骤。
    1. 创建算子工程。
    2. 实现算子,进行算子的原型定义、Kernel侧算子实现与Host侧Tiling实现。
    3. 算子入图开发,主要进行shape推导等算子入图适配函数的实现。
    4. 算子工程编译部署,编译生成自定义算子安装包并进行算子包的安装,将自定义算子部署到算子加速库。

    上述流程的详细描述可参见Ascend C算子开发指南中的算子入图(GE)图开发章节。

  1. Ascend Graph构图。

    基于自定义算子构造Ascend Graph图的方式与其他CANN内置算子相同,通过算子原型调用即可,典型流程可参见Ascend Graph构图