适配插件开发(PyTorch框架)
简介
本项目开发了NPU PyTorch算子插件,为使用PyTorch框架的开发者提供便捷的NPU算子库调用能力。 OP-Plugin算子插件的编译与使用均依赖昇腾PyTorch Adapter。在编译OpPlugin之前,请参见《CANN 软件安装指南》完成CANN软件的安装,参见《Ascend Extension for PyTorch 配置与安装》完成PyTorch框架的安装。本文档提供PyTorch适配算子开发指导,主要包括适配原则、适配文件结构和NPU适配算子开发三部分内容。
适配原则
- 不同的torch版本使用op_plugin_functions.yaml维护本身的对外接口。
- op_plugin对外接口与torch原生Aten IR保持一致。Aten IR接口说明,请参考pytorch/aten/src/ATen/native。
- 相同Aten IR算子存放于base_ops目录下,IR不同的算子存放于各自版本的目录下。不同适配方式的算子放置于不同的文件夹中,使用不同的命名空间,当前仅支持适配基于图IR执行算子(AclOp)和ACLNN算子(OpApi)。
- 非必要不使用NPUNativeFunction::命名空间中的接口。自定义算子使用custom_ops::xx调用,原生算子使用at::xx调用,调用其他适配接口使用op_plugin内部的接口,比如基于图IR执行算子使用acl_op::xx,ACLNN使用op_api::xx。
- 当前只支持基于图IR执行算子和ACLNN算子。
- 适配文件结构。
. ├── 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 ## 自动生成op_plugin对外接口的头文件,用于框架侧调用算子 │ ├── OpInterface.cpp ## 自动生成op_plugin对外接口路由实现,内部实现不同类型算子分支选择代码 │ ├── AclOpsInterface.h ## 自动生成基于图IR执行算子插件适配所对应头文件 │ ├── OpApiInterface.h ## 自动生成ACLNN算子插件适配所对应头文件 │ ├── ...
适配开发
PyTorch 1.11.0及以上版本官方提供的native_functions.yaml文件定义了PyTorch Native Functions的具体算子定义和分发细节,定义则通过.cpp文件实现。Op-Plugin仓库与原生类似,使用yaml文件定义了NPU适配的算子,算子具体适配则存放在.cpp文件中。
因此适配算子主要分为两步:
- 在yaml文件中配置算子的定义和其他配置。
- 需要完成算子适配的实现。
以torch API abs/abs_out为例,包含基于图IR执行算子和ACLNN算子,适配包括2部分,一是算子接口yaml配置,二是算子kernel的适配代码。
- 算子yaml配置。
Opplugin采用和原生PyTorch类似的逻辑在yaml中声明算子的各类信息,通过在yaml中配置算子,自动生成算子声明和注册代码。yaml文件的位置在op_plugin/config/内,不同的PyTorch版本有不同的子目录,如op_plugin/config/v2r1/op_plugin_functions.yaml表示PyTorch 2.1版本的算子配置文件目录。
yaml中算子配置规则如下面所示:
# 官方算子 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 #入参带有symint的算子 symint: - func: zeros(SymInt[] size, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor impl_ns: acl_op
文件说明参数说明:
- official表示该字段下的算子为官方原生;custom表示该字段下的算子为自定义算子;symint字段表明该算子支持symint类型的入参,该种算子可参考symint算子适配。
- func定义了算子的名称、入参和返回参数,具体规则可参考PyTorch中yaml的说明(LINK)。
- impl_ns:表示适配的算子类型,当前支持基于图IR执行算子acl_op和ACLNN算子op_api。如“impl_ns: acl_op,op_api”表明abs算子已有基于图IR执行算子适配实现和ACLNN适配实现。
- 算子适配实现。
当前支持适配基于图IR执行算子和ACLNN算子两类算子,当前算子适配文件存放于op_plugin/ops目录,其中如果算子在不同版本的PyTorch中实现一致,则存放于子目录base_ops内,根据算子类型存放于不同的子目录下。
- 基于图IR执行算子适配。
如版本间一致的abs算子的CANN适配文件路径为:op_plugin/ops/base_ops/aclops/AbsKernelNpu.cpp。
// 算子适配实现文件路径 op_plugin/ops/base_ops/aclops/AbsKernelNpu.cpp // 1. 引入依赖头文件 // 对外接口头文件,包含op_plugin所有基于图IR执行算子对外的函数原型 #include "op_plugin/AclOpsInterface.h" // torch调用基于图IR执行算子时,所依赖的基础函数对应的头文件 #include "op_plugin/utils/OpAdapter.h" // 2. 算子接口适配实现 // opplugin内适配的算子对外接口都定义在op_plugin命名空间中,外部调用方式为op_plugin::abs、op_plugin::abs_out;内部不同类型的算子适配采用不同的命名空间 // 基于图IR执行算子定义在acl_op命名空间中 namespace acl_op { using npu_preparation = at_npu::native::OpPreparation; using npu_utils = at_npu::native::NpuUtils; // 不对外暴露的接口,都定义在匿名空间中。常见为xx_nocheck等,直调基于图IR执行算子,不做内存、shape校验的函数。 namespace{ at::Tensor& abs_out_nocheck(at::Tensor& result, const at::Tensor& self) { at_npu::native::OpCommand cmd; cmd.Name("Abs") .Input(self) .Output(result) .Run(); return result; } } // namespace // abs_out api实现函数,参数与torch api一致。 at::Tensor& abs_out(const at::Tensor& self, at::Tensor& result) { // CheckOut作用:校验result的size、dtype等是否符合预期。若dtype不符合预期,则抛错。若size不符合则进行resize操作 npu_preparation::CheckOut({self}, result, self); // check_match作用:校验result是否为连续。因基于图IR执行算子无法支持非连续输出,result非连续时,需要单独处理。 if (!npu_utils::check_match(&result)) { // 若result非连续,创建连续tensor(contig_tensor),接收基于图IR执行算子(abs)的输出。再将contig_tensor拷贝到原始输出result。 at::Tensor contiguous_result = npu_utils::format_contiguous(result); abs_out_nocheck(contigTensor, self); npu_utils::format_fresh_view(result, contiguous_result); } else { // 若result连续,直接调用基于图IR执行算子。 abs_out_nocheck(result, self); } return result; } // abs api实现函数,参数与torch api一致。 at::Tensor abs(const at::Tensor& self) { // 构造输出tensor,调用基于图IR执行算子。 auto output_size = op_infer::infershape_for_elewise(self); at::Tensor result = npu_preparation::apply_tensor(self, output_size); abs_out_nocheck(result, self); return result; } // abs_ api实现函数,参数与torch api一致。该接口为inplace操作,即输出结果存放在输入tensor中。 at::Tensor& abs_(at::Tensor& self) { // 调用out接口,避免因self作为输出时,非连续场景下,直调基于图IR执行算子结果出错。 return acl_op::abs_out(self, self); return self; } } // namespace acl_op
- ACLNN算子适配。
如版本间一致的abs算子的ACLNN适配文件路径为:op_plugin/ops/base_ops/opApi/AbsKernelNpuOpApi.cpp。
//算子适配实现路径/op_plugin/ops/base_ops/opapi/AbsKernelNpuOpApi.cpp // 1. 引入依赖头文件 // 对外接口头文件,包含op_plugin所有ACLNN算子对外的函数原型 #include "op_plugin/OpApiInterface.h" // 引用 基于图IR执行算子头文件 #include "op_plugin/AclOpsInterface.h" // torch调用ACLNN算子时,所依赖的基础函数对应的头文件 #include "op_plugin/utils/op_api_common.h" // 2. 算子接口适配实现 // ACLNN算子定义在op_api命名空间中 namespace acl_op { using npu_preparation = at_npu::native::OpPreparation; // abs_out api实现函数,参数与torch api一致。 at::Tensor& abs_out(const at::Tensor& self, at::Tensor& result) { // 查找ACLNN算子实现,查找失败则使用基于图IR执行算子实现 DO_COMPATIBILITY(aclnnAbs, acl_op::abs_out(self, result)); npu_preparation::check_tensor({self}, result, self); // 异步调用npu执行 EXEC_NPU_CMD(aclnnAbs, self, result); return result; } // abs api实现函数,参数与torch api一致。 at::Tensor abs(const at::Tensor& self) { DO_COMPATIBILITY(aclnnAbs, acl_op::abs(self)); // construct the output tensor of the NPU at::Tensor result = npu_preparation::apply_tensor_without_format(self); // calculate the output result of the NPU EXEC_NPU_CMD(aclnnAbs, self, result); return result; } // abs_ api实现函数,参数与torch api一致。该接口为inplace操作,即输出结果存放在输入tensor中。 at::Tensor& abs_(at::Tensor& self) { DO_COMPATIBILITY(aclnnAbs, acl_op::abs_(self)); op_api::abs_out(self, self); return self; } } // namespace op_api
- 基于图IR执行算子适配。
自动前反向绑定算子配置
针对需要绑定前反向的算子(包括自定义算子和前反向绑定逻辑与原生不一致的原生算子)提供自动绑定前向算子和反向算子的功能。
- 适配前向和反向算子: 与上节算子适配开发中一致,分别适配前向算子和反向算子,并在op_plugin_functions.yaml中配置前向和反向算子。
- 配置前反向绑定,将前向和反向算子进行绑定: Opplugin与原生PyTorch一致,通过derivatives.yaml配置算子的前反向绑定关系,如下所示:
- name: l1_loss(Tensor self, Tensor target, int reduction=Mean) -> Tensor self: l1_loss_backward(grad, self, target, reduction)
symint算子适配
PyTorch在2.0版本及以上部分算子的入参变更为symint类型,在自动生成的代码中算子能够自动将symint类型的入参转换为int类型的入参,但部分情况需要直接适配symint类型的入参,此时要求:
- 算子在配置yaml中除了在official或custom下声明函数外,还需要同时在symint下配置该算子。
- 算子名称在原有名称上添加_symint后缀,如配置支持入参为symint类型的zeros算子,其yaml配置为:
# 官方算子 official: - func: zeros(SymInt[] size, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor impl_ns: acl_op symint: - func: zeros(SymInt[] size, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor impl_ns: acl_op
其算子实现如下,其中算子名称为zeros_symint,且入参中第一个阐述为symint相关的参数c10::SymIntArrayRef:
#include "op_plugin/AclOpsInterface.h" #include "op_plugin/utils/custom_functions/aclops/inner_compute.h" namespace acl_op { at::Tensor zeros_symint( c10::SymIntArrayRef size, c10::optional<at::ScalarType> dtype_opt, c10::optional<at::Layout> layout_opt, c10::optional<at::Device> device_opt, c10::optional<bool> pin_memory_opt) { return zeros_common_nocheck(c10::asIntArrayRefUnchecked(size), dtype_opt, layout_opt, device_opt, pin_memory_opt); } at::Tensor zeros( at::IntArrayRef size, c10::optional<at::DimnameList> names, c10::optional<at::ScalarType> dtype_opt, c10::optional<at::Layout> layout_opt, c10::optional<at::Device> device_opt, c10::optional<bool> pin_memory_opt) { return zeros_common_nocheck(size, dtype_opt, layout_opt, device_opt, pin_memory_opt); } } // namespace acl_op