下载
中文
注册
我要评分
文档获取效率
文档正确性
内容完整性
文档易理解
在线提单
论坛求助
昇腾小AI

单算子适配OpPlugin插件开发

简介

OpPlugin是Ascend Extension for PyTorch的算子插件,为使用PyTorch框架的开发者提供便捷的NPU算子库调用能力。OpPlugin算子插件的编译与使用均依赖昇腾Ascend Extension for PyTorch。在编译OpPlugin之前,请参见CANN 软件安装指南完成CANN软件的安装,参见Ascend Extension for PyTorch 配置与安装完成PyTorch框架的安装。本文档提供算子适配开发指导,主要包括适配原则、适配文件结构和算子适配开发三部分内容。

前提条件

如果用户使用快速安装方式安装torch_npu插件,适配前需执行如下命令拉取OP-Plugin仓对应分支的代码。
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文件中。

因此适配算子主要分为两步:

  1. 在yaml文件中配置算子。
  2. 完成算子适配的实现。

    以torch API abs/abs_out为例,包含基于aclop算子和aclnn算子,适配包括两部分,一是算子接口yaml配置,二是算子kernel的适配代码。

  1. 算子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
  2. 算子适配实现。

    当前支持适配基于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

自动前反向绑定算子配置

仅针对于需要进行前反向绑定的算子。

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插件。

  1. 编译生成二进制安装包。
    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。
  2. 完成编译后,安装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
搜索结果
找到“0”个结果

当前产品无相关内容

未找到相关内容,请尝试其他搜索词