适配插件开发（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

自动前反向绑定算子配置

针对需要绑定前反向的算子（包括自定义算子和前反向绑定逻辑与原生不一致的原生算子）提供自动绑定前向算子和反向算子的功能。

适配前向和反向算子：与上节算子适配开发中一致，分别适配前向算子和反向算子，并在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)
```
其中算子的配置规则可参考pytorch/aten/src/ATen/native。

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

父主题： 算子适配