文档HOT
下载
中文
注册

适配插件开发

本节将介绍自定义算子开发完成后，如何通过适配开发，从而在PyTorch框架中可以调用到该算子。

背景知识

通过PyTorch框架进行模型的训练、推理时，会调用到很多算子进行计算，目前PyTorch提供了常用的算子接口和接口的注册分发机制，可以将算子映射到不同的底层硬件设备。Ascend Extension for Pytorch中的OP-Plugin算子插件对此功能进行扩展，提供将PyTorch算子映射到昇腾AI处理器的功能。

PyTorch的适配流程，主要包括两个步骤：算子注册分发（yaml文件中配置算子的定义等）和适配插件实现。下文以PyTorch 1.11.0及以上官方版本为例，介绍PyTorch的适配流程。

单击LINK，您可以获取到OP-Plugin算子插件的更多信息并下载对应版本的源码包。需要开发者关注的适配文件目录结构如下：

.
├── op_plugin
│   ├── config                         # 算子配置文件目录
│   │   ├── derivatives.yaml          # 算子前反向绑定配置文件
│   │   └── op_plugin_functions.yaml  # 算子对外接口配置文件
│   ├── ops                            # 算子适配文件目录
│   │   ├── aclops                    # aclop算子适配目录
│   │   │   ├── AbsKernelNpu.cpp
│   │   │   └── ...
│   │   └── opapi                     # aclnn算子适配目录
│   │       ├── sparse                # sparse相关算子适配目录
│   │       │   └── SparseTensorUtils.h
│   │       ├── AbsKernelNpuOpApi.cpp
│   │       └── ...
│   ├── OpInterface.h         	  # 编译PyTorch框架后自动生成op_plugin对外接口的头文件,用于框架侧调用算子
│   ├── OpInterface.cpp               # 编译PyTorch框架后自动生成op_plugin对外接口路由实现，内部实现不同类型算子分支选择代码
│   ├── AclOpsInterface.h             # 编译PyTorch框架后自动生成ACLOP算子插件适配所对应头文件 
│   ├── OpApiInterface.h              # 编译PyTorch框架后自动生成ACLNN算子插件适配所对应头文件
│   ├── ...
  • 使用PyTorch框架完成自定义算子的调用，需要先完成算子的编译部署。
  • 编译部署时需要开启算子的二进制编译功能：修改算子工程中的编译配置项文件CMakePresets.json，将ENABLE_BINARY_PACKAGE设置为True。编译部署时可将算子的二进制部署到当前环境，便于后续算子的调用。
    "ENABLE_BINARY_PACKAGE": {
                    "type": "BOOL",
                    "value": "True"
                },
  • 已参考环境准备，配置CANN软件所需基本环境变量。
  • 编译部署后，执行PyTorch脚本前，需要将算子接口库的路径设置到共享库的查找路径中。下面示例仅作为参考，请根据实际情况进行设置。
    export LD_LIBRARY_PATH=$ASCEND_OPP_PATH/vendors/customize/op_api/lib/:$LD_LIBRARY_PATH

算子注册分发

对于自定义算子，由于没有具体的算子定义，我们需要在op_plugin_functions.yaml文件中给出定义，以便对算子进行结构化解析从而实现自动化注册和Python接口绑定。

op_plugin_functions.yaml文件介绍如下，您可以参考op_plugin_functions.yaml文件介绍，在文件中添加对应类型需要适配的算子信息。op_plugin_functions.yaml文件具体存放路径为：OP-Plugin算子插件源码包目录下的op_plugin/config/op_plugin_functions.yaml。

# op_plugin_functions.yaml
all_version: [v1.11, v2.0, v2.1, v2.2, v2.3, v2.4, v2.5]
# 官方算子
official:
  - func: abs(Tensor self) -> Tensor
    acl_op: all_version
    op_api: v1.11, v2.1, v2.2, v2.3, v2.4, v2.5
    gen_opapi:
      structured_inherit: abs.out
# 自定义算子
custom:
  - func: my_abs(Tensor self) -> Tensor
    acl_op: v1.11, v2.1, v2.2, v2.3, v2.4, v2.5
    op_api: all_version

文件说明参数说明：

  • all_version表示当前PyTorch支持的所有版本。
  • official和custom分别表示该字段下的算子为原生和自定义算子。
  • func定义了算子的schema，主要有名称、入参和返回参数，具体规则可参考原生定义（LINK）。
  • acl_op字段后面填版本名称，表示在该版本支持acl_op调用，如果支持的版本与all_version表示的版本一致，则可以用"all_version"表示，可选字段。
  • op_api字段后面填版本名称，表示在该版本支持op_api调用，如果支持的版本与all_version表示的版本一致，则可以用"all_version"表示，可选字段。
  • gen_opapi字段此版本暂未开放使用。

    如果存在某个Aten IR有两个版本不一致，则需要两个都加上，如std.correction在1.11和2.1及以上的入参名称不同，则需要分开写成两个，通过version区分。

      - func: std.correction(Tensor self, int[1]? dim, *, int? correction, bool keepdim=False) -> Tensor
    acl_op: v1.11
    op_api: v1.11

  - func: std.correction(Tensor self, int[1]? dim=None, *, Scalar? correction=None, bool keepdim=False) -> Tensor
    acl_op: v2.1, v2.2, v2.3, v2.4, v2.5
    op_api: v2.1, v2.2, v2.3, v2.4, v2.5

适配插件开发

下面介绍开发PyTorch适配的具体过程。用户通过开发算子适配插件，实现PyTorch原生算子的输入参数、输出参数和属性的格式转换，使转换后的格式与自定义算子的输入参数、输出参数和属性的格式相同。

  1. 创建适配插件文件。

    Ascend C算子适配文件保存路径请参考背景知识中的目录结构说明，命名风格采用大驼峰，命名格式：<算子名> + <KernelNpu>.cpp，如：AddCustomKernelNpu.cpp。

  2. 引入依赖头文件。
    // 对外接口头文件，包含op_plugin所有ACLNN算子对外的函数原型
#include "op_plugin/OpApiInterface.h"
// 引用 基于图IR执行算子头文件
#include "op_plugin/AclOpsInterface.h"
// torch调用ACLNN算子时，所依赖的基础函数对应的头文件
#include "op_plugin/utils/op_api_common.h"
  3. 定义实现算子适配主体函数。

    实现算子适配主体函数，根据Ascend C算子原型构造得到对应的input、output、attr。

  4. 重编译PyTorch框架或插件。

    请重新编译生成torch_npu插件安装包并安装。

父主题： PyTorch框架