概述
功能简介
默认情况下,昇腾NPU的自定义算子无法正常入图PyTorch,为了能够以图模式使能这些自定义算子,TorchAir提供了一种插件化入图方式。该方式借助PyTorch原生注册机制,无需重新编译Ascend Extension for PyTorch的torch_npu和TorchAir就能完成算子入图,方便用户进行后续算子调试或将自定义算子模块作为插件使用。
目前昇腾提供了基于Ascend C自定义算子开发方式,本章将详细阐述如何以插件化方式实现自定义NPU算子入图,入图流程具体参见图1。
使用方法
- 环境准备。
参考安装,完成配套版本的CANN、PyTorch、torch_npu安装。
- 算子开发和部署。
基于Ascend C进行自定义算子开发,完成算子入GE图(如提供shape推导函数等),并完成算子编译和部署。
- PyTorch侧适配。
- 算子注册PyTorch。
- 算子converter注册和实现。
- 通过register_fx_node_ge_converter接口将Aten算子注册到TorchAir。
- 通过custom_op接口实现自定义算子的converter,完成Aten IR到GE IR的转换。
父主题: 自定义算子插件化入图