单算子适配OpPlugin插件开发
简介
OpPlugin是Ascend Extension for PyTorch的算子插件,为使用PyTorch框架的开发者提供便捷的NPU算子库调用能力。OpPlugin算子插件的编译与使用均依赖昇腾Ascend Extension for PyTorch。在编译OpPlugin之前,请参见《CANN 软件安装指南》完成CANN软件的安装,参见《Ascend Extension for PyTorch 配置与安装》完成PyTorch框架的安装。本文档提供算子适配开发指导,主要包括适配原则、适配文件结构和算子适配开发三部分内容。
前提条件
git clone https://gitee.com/ascend/op-plugin.git -b 6.0.rc3 cd op-plugin
适配原则
- OpPlugin对外接口与torch原生Aten IR保持一致。Aten IR接口说明,请参考pytorch/aten/src/ATen/native。
- 不同的torch版本使用op_plugin_functions.yaml维护本身的对外接口,配置规则可参见算子适配开发。
- 不同适配方式的算子放置于不同的文件夹中,使用不同的命名空间,当前仅有aclop适配算子(通过GE注册的算子)和opapi适配算子(aclnn等)。
- 非必要不使用NPUNativeFunction::命名空间中的接口。自定义算子使用custom_ops::xx调用,原生算子使用at::xx调用,调用其他适配接口使用OpPlugin内部的接口,比如aclop算子使用acl_op::xx,aclnn使用op_api::xx。
- 当前只支持aclop算子和aclnn算子。
适配文件结构
. ├── 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官方提供的native_functions.yaml文件定义了PyTorch Native Functions的具体算子定义和分发细节,定义则通过.cpp文件实现。OpPlugin仓库与原生类似,使用yaml文件定义了NPU适配的算子,算子具体适配则存放在.cpp文件中。
因此适配算子主要分为两步:
- 在yaml文件中配置算子。
- 完成算子适配的实现。
以torch API abs/abs_out为例,包含基于aclop算子和aclnn算子,适配包括两部分,一是算子接口yaml配置,二是算子kernel的适配代码。
- 算子yaml配置。
OpPlugin采用和原生PyTorch类似的逻辑在yaml中声明算子的各类信息,通过在yaml中配置算子,自动生成算子声明和注册代码。算子的Aten IR定义位于op_plugin/config/op_plugin_functions.yaml文件中,所有版本的定义都在这个文件里面,通过配置不同版本来区分。
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 #入参带有symint的算子 symint: - func: zeros(SymInt[] size, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor acl_op: v2.1, v2.2, v2.3, v2.4, v2.5
文件说明参数说明:
- all_version表示当前PyTorch支持的所有版本。
- official和custom分别表示该字段下的算子为原生和自定义算子;symint字段表明该算子支持symint类型的入参,该种算子请参考symint算子适配。
- func定义了算子的schema,主要有名称、入参和返回参数,具体规则可参考原生定义(LINK)。
- acl_op字段后面填版本名称,表示在该版本支持acl_op调用,如果支持的版本与all_version表示的版本一致,则可以用"all_version"表示,可选字段。
- op_api字段后面填版本名称,表示在该版本支持op_api调用,如果支持的版本与all_version表示的版本一致,则可以用"all_version"表示,可选字段。
- gen_opapi对于支持op_api调用的算子,如果适配代码简单,可以直接调用底层算子,不需要额外的适配,则可以考虑用结构化适配的方式自动生成适配代码,详见章节结构化适配介绍(可选)。
如果存在某个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
- 算子适配实现。
当前支持适配基于aclop算子和aclnn算子两类算子,aclop算子适配文件位于op_plugin/ops/aclops目录,aclnn算子适配文件位于op_plugin/ops/opapi目录。一个算子所有版本的适配代码都在一个文件中,通过编译宏VERSION_BETWEEN来区分不同版本。
- aclop算子适配。
如果所有版本的适配代码一致,则不需要额外添加编译宏,适配文件路径为:op_plugin/ops/aclops/AbsKernelNpu.cpp,文件命名规范为算子名称+KernelNpu,算子名称首字母大写。
// 算子适配实现文件路径 op_plugin/ops/aclops/AbsKernelNpu.cpp // 1. 引入依赖头文件 // 对外接口头文件,包含op_plugin所有aclop算子对外的函数原型 #include "op_plugin/AclOpsInterface.h" // torch调用ACLOP算子时,所依赖的基础函数对应的头文件 #include "op_plugin/utils/OpAdapter.h" // 2. 算子接口适配实现 // opplugin内适配的算子对外接口都定义在op_plugin命名空间中,外部调用方式为op_plugin::abs、op_plugin::abs_out;内部不同类型的算子适配采用不同的命名空间 // CANN算子定义在acl_op命名空间中, namespace acl_op { using npu_preparation = at_npu::native::OpPreparation; using npu_utils = at_npu::native::NpuUtils; // 不对外暴露的接口,都定义在匿名空间中。常见为xx_nocheck等,直调ACLOP算子,不做内存、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是否为连续。因ACLOP算子无法支持非连续输出,result非连续时,需要单独处理。 if (!npu_utils::check_match(&result)) { // 若result非连续,创建连续tensor(contig_tensor),接收ACLOP算子(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连续,直接调用ACLOP算子。 abs_out_nocheck(result, self); } return result; } // abs api实现函数,参数与torch api一致。 at::Tensor abs(const at::Tensor& self) { // 构造输出tensor,调用ACLOP算子。 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作为输出时,非连续场景下,直调ACLOP算子结果出错。 return acl_op::abs_out(self, self); return self; } } // namespace acl_op
不同版本间适配代码有差异的,所有代码均放在同一个文件中,用编译宏来区分。
#include "op_plugin/AclOpsInterface.h" #include "op_plugin/utils/custom_functions/aclops/inner_compute.h" namespace acl_op { // 1.11的函数入参和2.0及以上版本有区别,因此用宏来控制 #if VERSION_BETWEEN(V1R11, V1R11) at::Tensor embedding( const at::Tensor& weight, const at::Tensor& indices, int64_t padding_idx, bool scale_grad_by_freq, bool sparse) { return embedding_common_nocheck(weight, indices); } #endif // 2.0及以上版本的代码都一致 #if VERSION_BETWEEN(V2R0, VERSION_NEWEST) at::Tensor embedding_symint( const at::Tensor& weight, const at::Tensor& indices, c10::SymInt padding_idx, bool scale_grad_by_freq, bool sparse) { return embedding_common_nocheck(weight, indices); } #endif } // namespace acl_op
- aclnn算子适配。
aclnn算子适配与aclop类似,也是如果所有版本的适配代码一致,则不需要额外添加编译宏,适配文件路径为:op_plugin/ops/opapi/AbsKernelNpuOpApi.cpp,文件命名规范为算子名称+KernelNpuOpApi,算子名称首字母大写。
//算子适配实现路径/op_plugin/ops/base_ops/opapi/AbsKernelNpuOpApi.cpp // 1. 引入依赖头文件 // 对外接口头文件,包含op_plugin所有ACLNN算子对外的函数原型 #include "op_plugin/OpApiInterface.h" // 引用 ACLOP 算子声明头文件 #include "op_plugin/AclOpsInterface.h" // torch调用ACLNN算子时,所依赖的基础函数对应的头文件 #include "op_plugin/utils/op_api_common.h" // 2. 算子接口适配实现 // ACLNN算子定义在op_api命名空间中, namespace op_api { using npu_preparation = at_npu::native::OpPreparation; // abs_out api实现函数,参数与torch api一致。 at::Tensor& abs_out(const at::Tensor& self, at::Tensor& result) { // 查找ACLNN算子实现,查找失败则使用ACLOP算子实现 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操作,即输出结果存放在输入 at::Tensor& abs_(at::Tensor& self) { DO_COMPATIBILITY(aclnnAbs, acl_op::abs_(self)); op_api::abs_out(self, self); return self; } } // namespace op_api
不同版本间适配代码有差异的,所有代码均放在同一个文件中,用编译宏来区分。
#include "op_plugin/AclOpsInterface.h" #include "op_plugin/OpApiInterface.h" #include "op_plugin/utils/op_api_common.h" namespace op_api { using npu_preparation = at_npu::native::OpPreparation; // 1.11的函数入参和2.0及以上版本有区别,需要单独实现,因此用宏来控制 #if VERSION_BETWEEN(V1R11, V1R11) at::Tensor embedding(const at::Tensor& weight, const at::Tensor& indices, int64_t padding_idx, bool scale_grad_by_freq, bool sparse) { DO_COMPATIBILITY(aclnnEmbedding, acl_op::embedding(weight, indices, padding_idx, scale_grad_by_freq, sparse)); // calculate the output size auto output_size = op_infer::array_to_small_vector(indices.sizes()); output_size.emplace_back(weight.size(weight.dim() - 1)); // construct the output tensor of the NPU at::Tensor result = npu_preparation::apply_tensor_without_format(output_size, weight.options()); // calculate the output resugt of the NPU EXEC_NPU_CMD(aclnnEmbedding, weight, indices, result); return result; } #endif #if VERSION_BETWEEN(V2R0, VERSION_NEWEST) at::Tensor embedding_symint( const at::Tensor& weight, const at::Tensor& indices, c10::SymInt padding_idx, bool scale_grad_by_freq, bool sparse) { DO_COMPATIBILITY(aclnnEmbedding, acl_op::embedding_symint(weight, indices, padding_idx, scale_grad_by_freq, sparse)); // calculate the output size auto output_size = op_infer::array_to_small_vector(indices.sizes()); output_size.emplace_back(weight.size(weight.dim() - 1)); // construct the output tensor of the NPU at::Tensor result = npu_preparation::apply_tensor_without_format(output_size, weight.options()); // calculate the output resugt of the NPU EXEC_NPU_CMD(aclnnEmbedding, weight, indices, result); return result; } #endif } // namespace op_api
- aclop算子适配。
自动前反向绑定算子配置
仅针对于需要进行前反向绑定的算子。
PyTorch的算子自动反向微分依赖于算子的前反向绑定,即前向函数和反向函数的绑定。对于原生的算子,官方已有前反向绑定逻辑,插件侧有对应前向算子和反向算子适配即可(只需要在op_plugin_functions.yaml里面配置)。对于自定义算子,则需要在插件侧配置前反向自动绑定。
针对需要绑定前反向的算子(包括自定义算子和前反向绑定逻辑与原生不一致的原生算子)提供自动绑定前向算子和反向算子的功能。
- 适配前向和反向算子: 与算子适配开发中一致,分别适配前向算子和反向算子,并在op_plugin_functions.yaml中配置前向和反向算子。
- 配置前反向绑定,将前向和反向算子进行绑定: OpPlugin与原生PyTorch一致,通过derivatives.yaml配置算子的前反向绑定关系,同时相比原生新增了version字段用于表示支持的版本,如下所示:
# derivatives.yaml - name: l1_loss(Tensor self, Tensor target, int reduction=Mean) -> Tensor self: l1_loss_backward(grad, self, target, reduction) target: l1_loss_backward(grad, self, target, reduction) * -1 version: v2.0, v2.1, v2.2, v2.3, v2.4, v2.5
所有版本的算子前反向绑定都在同一个derivatives.yaml里面,通过version字段来区分版本。
symint算子适配
Symint类型算子需参考此部分进行适配。
Symint为PyTorch在2.0及以上版本新增的数据类型,yaml配置中对应添加了symint字段。配置在symint字段下的函数表示底层函数实现支持了Symint类型入参。对于底层不支持Symint的函数,则无需在Symint字段配置。部分情况需要在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 acl_op: v2.1, v2.2, v2.3, v2.4, v2.5 symint: - func: zeros(SymInt[] size, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor acl_op: v2.1, v2.2, v2.3, v2.4, v2.5
其算子实现如下,其中算子名称为zeros_symint,且入参中第一个参数的类型为symint相关的类型c10::SymIntArrayRef,同时由于symint特性只在Pytorch2.0以上支持,symint相关适配代码需要根据实际版本支持情况添加版本编译宏VERSION_BETWEEN来控制编译:
#include "op_plugin/AclOpsInterface.h" #include "op_plugin/utils/custom_functions/aclops/inner_compute.h" namespace acl_op { #if VERSION_BETWEEN(V2R1, VERSION_NEWEST) 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); } #endif 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
结构化适配介绍(可选)
仅对支持op_api的算子可使用此方法进行适配。
结构化适配指通过在op_plugin_functions.yaml中进行配置,可自动生成算子实现Kernel。判断是否可结构化依据:opapi对应的aclnn算子与Aten IR的语义对齐,适配层除申请output tensor,无其他适配逻辑。自动生成的适配文件位于op_plugin/ops/opapi/StructKernelNpuOpApi.cpp。
Yaml配置有以下两种方式,可根据实际情况进行选择。每个结构化适配的函数必须在op_plugin_functions.yaml中配置,具有如下格式:
- 方式一(常规场景):
- func: func_name(ArgType arg0[=default], ArgType arg1[=default], ...) -> Return op_api: v2.1 gen_opapi: out: size: arg0 dtype: arg1.scalar_type() name: arg0 exec: aclnnFuncName
各个字段的含义如下:
- gen_opapi: 表示对应API可结构化,其他字段需要配置在此字段下。
- out: 表示函数的输出,此字段下面包含size和dtype字段,如果包含多个输出,可配置成out0、out1等。对于out类接口,此字段不可自定义,需要与Aten IR定义的输出参数名相同。对于inplace类接口,不需要配置此字段。
- size: 配置输出tensor的shape大小,如果大小和schema中的某个参数相同,可以配置成输入参数的名字。也可配置成自定义infershape函数,infershape函数需在KernelNpuOutputSize.h中实现。对于out类接口,如果输出shape不变,可省略此字段。配置方式主要包含以下几种:
- func: func_name(ArgType arg0[=default], ArgType arg1[=default], ...) -> Return Aten IR定义: - func: func_name(ArgType arg0, ArgType arg1, ...) -> Return 方式一: 和输入参数相同 size: arg0 方式二:枚举每个维度的值 size: '{4, arg0.size(0), arg0.size(1), arg1.size(0)}' 方式三:条件表达式 size: 'arg1 == 1? arg0.sizes(): at::ArrayRef<int64_t>()' 方式四:在KernelNpuOutputSize.h中自定义infershape函数, 例如broadcast_ops_npu_output_size size: broadcast_ops_npu_output_size(arg0, arg1)
- dtype: 配置输出tensor的dtype大小,如果大小和schema中的某个参数相同,可以配置成输入参数的名字。也可配置成自定义inferdtype函数,inferdtype函数需在KernelNpuOutputDtype.h中实现。对于out类接口,如果输出dtype不需要check,可省略此字段。配置方式主要包含以下几种:
Aten IR定义: - func: func_name(ArgType arg0, ArgType arg1, ...) -> Return 方式一: 和输入参数相同 dtype: arg0 方式二:配置成已知的dtype类型 dtype: at::kFloat 方式三:条件表达式 dtype: 'isIntegralType(arg0.scalar_type(), true) ? at::kFloat : arg0.scalar_type()' 方式四:在KernelNpuOutputDtype.h中自定义inferdtype函数。 dtype: inferdtype(arg0, arg1)
- name: 输出结果涉及named tensor逻辑,可配置此字段,当前仅支持name和输入参数相同的配置,不涉及可忽略。
- exec: 配置EXEC_NPU_CMD对应的参数,如果除aclnnname,其它参数顺序和Aten IR的顺序相同,可只配置aclnnname,如aclnnAbs。以abs为例,exec字段可以配置成下面两种方式。
- func: abs.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) 方式一: exec: aclnnAbs, self, out 方式二: exec: aclnnAbs
- 方式二(继承场景):
- func: func_name(ArgType arg0[=default], ArgType arg1[=default], ...) -> Return op_api: v2.1 gen_opapi: structured_inherit: func_name.out
structured_inherit:如果原函数或inplace类接口的字段配置与out类接口的字段配置相同,可通过此字段继承对应的out类接口。
以abs为例,原函数和out类函数的out属性和exec相同,可通过structured_inherit字段继承。- func: abs(Tensor self) -> Tensor op_api: v1.11, v2.1, v2.2, v2.3, v2.4, v2.5 gen_opapi: structured_inherit: abs.out
后续处理
算子适配完成后,需按如下操作重新编译OpPlugin插件。
- 编译生成二进制安装包。
cd op-plugin bash ci/build.sh --python=3.8 --pytorch=v2.1.0-6.0.rc3
- 指定Python版本编包方式,以Python3.8为例,其他Python版本请使用 --python=3.9或--python3.10。
- 指定PyTorch版本编包方式,以PyTorch2.1.0为例,其他PyTorch版本请使用--pytorch=v2.3.1或--pytorch=v2.4.0。
- 指定Ascend Extension for PyTorch软件版本6.0.rc3。
- 完成编译后,安装dist目录下生成的插件torch_npu包,如果使用非root用户安装,需要在命令后加--user。
# 请用户根据实际情况更改命令中的torch_npu包名 pip3 install --upgrade dist/torch_npu-2.1.0.post8-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl