单算子适配OpPlugin插件开发

简介

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

前提条件

如果用户使用快速安装方式安装torch_npu插件，适配前需执行如下命令拉取torch_npu仓对应分支的代码并进入OpPlugin目录。
git clone https://gitee.com/ascend/pytorch.git -b v2.1.0-6.0.0 --recursive
cd pytorch/third_party/op-plugin

适配原则

适配文件结构

.
├── 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文件中。

以下abs的yaml配置和适配文件为已有配置和文件，此处仅为示例，用户需根据实际场景更改。

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

  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中算子配置规则如下面所示：

     1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
    # 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区分。

      1
2
3
4
5
6
7
        - 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，算子名称首字母大写。

       1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
      // 算子适配实现文件路径 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);
}
} // namespace acl_op

      不同版本间适配代码有差异的，所有代码均放在同一个文件中，用编译宏来区分。

       1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
      #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，算子名称首字母大写。

       1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
      //算子适配实现路径 op_plugin/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

      不同版本间适配代码有差异的，所有代码均放在同一个文件中，用编译宏来区分。

       1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
      #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 result 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 result of the NPU
    EXEC_NPU_CMD(aclnnEmbedding, weight, indices, result);
    return result;
}
#endif

} // namespace op_api

自动前反向绑定算子配置

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

PyTorch的算子自动反向微分依赖于算子的前反向绑定，即前向函数和反向函数的绑定。对于原生的算子，官方已有前反向绑定逻辑，插件侧有对应前向算子和反向算子适配即可(只需要在op_plugin_functions.yaml里面配置)。对于自定义算子，则需要在插件侧配置前反向自动绑定。

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

symint算子适配

Symint类型算子需参考此部分进行适配。

以下yaml配置和适配文件为已有配置和文件，此处仅为示例，用户需根据实际场景更改。

Symint为PyTorch在2.0及以上版本新增的数据类型，yaml配置中对应添加了symint字段。配置在symint字段下的函数表示底层函数实现支持了Symint类型入参。对于底层不支持Symint的函数，则无需在Symint字段配置。部分情况需要在Symint字段配置时，用户需进行如下操作进行算子适配:

结构化适配介绍（可选）

仅对支持op_api的算子可使用此方法进行适配。

结构化适配指通过在op_plugin_functions.yaml中进行配置，可自动生成算子实现Kernel。判断是否可结构化依据：opapi对应的aclnn算子与Aten IR的语义对齐，适配层除申请output tensor，无其他适配逻辑。自动生成的适配文件位于op_plugin/ops/opapi/StructKernelNpuOpApi.cpp。

Yaml配置有以下两种方式，可根据实际情况进行选择。每个结构化适配的函数必须在op_plugin_functions.yaml中配置，具有如下格式：

后续处理

算子适配完成后，需按如下操作重新编译torch_npu包。

  1. 编译生成二进制安装包。
    # 回到torch_npu目录
cd ../../
bash ci/build.sh --python=3.8

    指定Python版本编包方式，以Python3.8为例，其他Python版本请使用 --python=3.9或--python3.10。

  2. 完成编译后，安装dist目录下生成的插件torch_npu包，如果使用非root用户安装，需要在命令后加--user。
    # 请用户根据实际情况更改命令中的torch_npu包名
pip3 install --upgrade dist/torch_npu-2.1.0.post10-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
父主题： 昇腾自研插件