适配插件开发（PyTorch框架）

简介

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

适配原则

适配开发

PyTorch 1.11.0及以上版本官方提供的native_functions.yaml文件定义了PyTorch Native Functions的具体算子定义和分发细节，定义则通过.cpp文件实现。Op-Plugin仓库与原生类似，使用yaml文件定义了NPU适配的算子，算子具体适配则存放在.cpp文件中。

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

  1. 在yaml文件中配置算子的定义和其他配置。
  2. 需要完成算子适配的实现。

    以torch API abs/abs_out为例，包含基于图IR执行算子和ACLNN算子，适配包括2部分，一是算子接口yaml配置，二是算子kernel的适配代码。

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

  2. 算子适配实现。

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

自动前反向绑定算子配置

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

symint算子适配

PyTorch在2.0版本及以上部分算子的入参变更为symint类型，在自动生成的代码中算子能够自动将symint类型的入参转换为int类型的入参，但部分情况需要直接适配symint类型的入参，此时要求：

父主题： 算子适配