适配插件开发(PyTorch框架)
本节将介绍自定义算子开发完成后,如何通过适配开发,从而在PyTorch框架中可以调用到该算子。
背景知识
通过PyTorch框架进行模型的训练、推理时,会调用到很多算子进行计算,目前PyTorch提供了常用的算子接口和接口的注册分发机制,可以将算子映射到不同的底层硬件设备。Ascend Extension for Pytorch中的OP-Plugin算子插件对此功能进行扩展,提供将PyTorch算子映射到昇腾AI处理器的功能。
PyTorch的适配流程,主要包括两个步骤:算子注册分发(yaml文件中配置算子的定义等)和适配插件实现。下文以PyTorch 1.11.0及以上官方版本为例,介绍PyTorch的适配流程。
. ├── op_plugin │ ├── config # 算子配置文件目录 │ │ ├── v2r1 # pt2.1版本算子配置文件目录 │ │ │ ├── op_plugin_functions.yaml # 算子对外接口配置文件 │ │ │ ├── derivatives.yaml # 算子前反向绑定的配置文件 │ │ │ ... │ ├── ops │ │ ├── base_ops # 多版本适配一致的算子目录 │ │ │ ├── aclops # 基于图IR执行算子适配文件目录 │ │ │ │ ├── abs.cpp # abs算子适配文件 │ │ │ ├── opapi # ACLNN适配目录 │ │ │ │ ├── absOpApi.cpp # abs算子适配文件 │ │ ├── v2r1 # 2.1版本适配目录 │ │ │ │ ├── aclops │ │ │ │ ├── opapi │ │ ├──... │ ├── OpInterface.h # 编译PyTorch框架后自动生成op_plugin对外接口的头文件,用于框架侧调用算子 │ ├── OpInterface.cpp # 编译PyTorch框架后自动生成op_plugin对外接口路由实现,内部实现不同类型算子分支选择代码 │ ├── AclOpsInterface.h # 编译PyTorch框架后自动生成基于图IR执行算子插件适配所对应头文件 │ ├── 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/xxxx/op_plugin_functions.yaml(xxxx代表对应的PyTorch版本)。
# 官方算子 official: - func: abs(Tensor self) -> Tensor impl_ns: acl_op, op_api - func: zeros(SymInt[] size, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor impl_ns: acl_op # 自定义算子 custom: - func: my_abs(Tensor self) -> Tensor impl_ns: acl_op
参数说明:
适配插件开发
下面介绍开发PyTorch适配的具体过程。用户通过开发算子适配插件,实现PyTorch原生算子的输入参数、输出参数和属性的格式转换,使转换后的格式与自定义算子的输入参数、输出参数和属性的格式相同。
- 创建适配插件文件。
Ascend C算子适配文件保存路径请参考背景知识中的目录结构说明,命名风格采用大驼峰,命名格式:<算子名> + <KernelNpu>.cpp,如:AddCustomKernelNpu.cpp。
- 引入依赖头文件。
// 对外接口头文件,包含op_plugin所有ACLNN算子对外的函数原型 #include "op_plugin/OpApiInterface.h" // 引用 基于图IR执行算子头文件 #include "op_plugin/AclOpsInterface.h" // torch调用ACLNN算子时,所依赖的基础函数对应的头文件 #include "op_plugin/utils/op_api_common.h"
- 定义实现算子适配主体函数。
实现算子适配主体函数,根据Ascend C算子原型构造得到对应的input、output、attr。
- 重编译PyTorch框架或插件。
请重新编译生成torch_npu插件安装包并安装。