昇腾社区首页
中文
注册
开发者
下载

注册并适配Eager模式

完成算子NPU实现后,可对接PyTorch的Eager模式进行适配。Ascend Extension for PyTorch提供了OpPlugin算子插件,用来实现PyTorch算子注册Eager模式适配

本章仅提供OpPlugin适配的关键步骤说明,详细的操作请参考《PyTorch 框架特性指南》中的“基于OpPlugin算子适配开发”章节,例如算子yaml配置、算子适配等实现。

  1. 注册PyTorch算子
  2. 基于OpPlugin适配Eager模式

注册PyTorch算子

PyTorch官方提供的native_functions.yaml文件定义了PyTorch Native Functions的具体算子定义和分发细节,定义则通过*.cpp文件实现。OpPlugin库与原生库类似,也使用yaml文件定义了NPU适配的算子,算子具体适配则存放在*.cpp文件中。

请确保已按环境准备下载torch_npu源码,算子的ATen IR定义位于third_party/op-plugin/op_plugin/config/op_plugin_functions.yaml文件中,在“custom字段”下添加目标PyTorch算子Schema定义:

custom:   
- func: my_inplace(Tensor(a!) x, Tensor(b!) y) -> ()
  op_api: all_version

上述原型定义对应的PyTorch算子为torch.ops.npu.my_inplace,其中torch.ops是PyTorch算子固定开头,npu为torch_npu自定义算子库的名称,my_inplace为自定义算子名。

基于OpPlugin适配Eager模式

完成算子NPU实现(Ascend C)和PyTorch算子注册后,需要在PyTorch的Eager模式适配层调用Ascend C算子。

借助OpPlugin插件提供的工程化适配能力,简化Eager模式适配层开发,在third_party/op-plugin/op_plugin/config/op_plugin_functions.yaml的my_inplace算子原型注册下,追加“gen_opapi”字段(表示对应可结构化的API):

- func: my_inplace(Tensor(a!) x, Tensor(b!) y) -> ()
  op_api: all_version
  gen_opapi:
    exec: aclnnMyInplace      # 等价于aclnnMyInplace,x,y

exec:配置对应的EXEC_NPU_CMD接口,一般指aclnnXxx前缀接口。本样例配置为aclnnMyInplace,aclnn为固定前缀,MyInplace为Ascend C算子名,表示调用Ascend C算子MyInplace实现PyTorch算子my_inplace。