下载
中文
注册
我要评分
文档获取效率
文档正确性
内容完整性
文档易理解
在线提单
论坛求助
昇腾小AI

概述

功能简介

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

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

图1 自定义算子入图流程

使用方法

  1. 环境准备。

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

  2. 算子开发和部署

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

  3. PyTorch侧适配。
    1. 算子注册PyTorch
      1. 通过OP-Plugin完成算子在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算子工程。

搜索结果
找到“0”个结果

当前产品无相关内容

未找到相关内容,请尝试其他搜索词