概述

功能简介

默认情况下，昇腾NPU的自定义算子无法正常入图PyTorch，为了能够以图模式使能这些自定义算子，TorchAir提供了一种插件化入图方式。该方式借助PyTorch原生注册机制，无需重新编译Ascend Extension for PyTorch的torch_npu和TorchAir就能完成算子入图，方便用户进行后续算子调试或将自定义算子模块作为插件使用。

目前昇腾提供了基于Ascend C自定义算子开发方式，本章将详细阐述如何以插件化方式实现自定义NPU算子入图，入图流程具体参见图1。

1. 环境准备。

   参考安装，完成配套版本的CANN、PyTorch、torch_npu安装。

2. 算子开发和部署。

   基于Ascend C进行自定义算子开发，完成算子入GE图（如提供shape推导函数等），并完成算子编译和部署。

3. PyTorch侧适配。
   1. 算子注册PyTorch。
      1. 通过OpPlugin完成算子在PyTorch注册，并提供torch_npu运行所需的算子适配文件。
      2. 在PyTorch注册Meta函数，以完成入图时的shape与dtype推导。
   2. 算子converter注册和实现。
      1. 通过register_fx_node_ge_converter接口将Aten算子注册到TorchAir。
      2. 通过custom_op接口实现自定义算子的converter，完成Aten IR到GE IR的转换。

样例参考

本文档仅提供自定义算子入图的关键步骤指导，端到端的算子入图、运行过程请单击Gitee进入Ascend samples仓，获取自定义AddCustom算子入图样例代码。

请仔细阅读README.md，了解PyTorch如何以图模式调用AddCustom算子工程。