适配插件开发（TensorFlow框架）

简介

您可以参考本章节进行算子适配插件的开发，将第三方框架的算子映射成适配昇腾AI处理器的算子。基于TensorFlow 框架的网络运行时，首先会加载并调用Graph Engine（简称GE）中的插件信息，将原始框架网络中的算子进行解析并映射成适配昇腾AI处理器算子。

下文我们将适配昇腾AI处理器的算子称为CANN算子。

原理介绍

算子插件的实现包含CANN算子类型的注册、原始框架中算子类型的注册以及原始框架中算子属性到CANN算子属性的映射，算子的映射通过Parser模块完成。插件在整网络运行场景下的实现流程如图1所示。

图1 算子插件的实现流程

1. 首先GE接收到第三方框架的原始网络模型，并进行初始化，网络模型的拓扑图我们简称为图。
2. GE从Register注册模块中加载算子插件生成的.so文件，在CANN软件安装后文件存储路径的“opp/framework/”路径中。
3. 读取算子插件.so中的算子相关信息，并将其注册到算子插件的map文件中（所有算子插件的相关信息都会以map的形式存储到一个文件中）。
4. GE向Parser模块发送调用Parser方法的请求。
5. Parser模块根据算子类型（OpType）从算子插件的map文件中取出对应的Parser函数，并返回实现函数ParseParamsByOperatorFn给Parser模块，Parser模块根据实现函数将第三方网络算子中的属性映射到CANN算子的属性，即算子原型中的属性定义，从而完成第三方网络中算子到CANN算子的映射。
6. 后续会进行图准备、图拆分及图优化等一系列操作，最终生成适配昇腾AI处理器的网络模型。

插件实现

GE提供REGISTER_CUSTOM_OP宏，按照指定的算子名称完成算子的注册。

#include "register/register.h"
#include "graph/operator.h"
namespace domi
{
REGISTER_CUSTOM_OP("OpType")
    .FrameworkType(TENSORFLOW) 
    .OriginOpType("OriginOpType")
    .ParseParamsByOperatorFn(ParseParamByOpFunc)
    .ImplyType(ImplyType::TVM);      // TBE算子：ImplyType::TVM；AI CPU算子：ImplyType::AI_CPU
}

下面详细介绍解析函数ParseParamsByOperatorFn，此函数的实现有以下几种场景：

.ParseParamsByOperatorFn(AutoMappingByOpFn) 

• 算子属性无法一一对应

  例如TensorFlow中的TopKV2定义如下所示：

  REGISTER_OP("TopKV2")
      .Input("input: T")
      .Input("k: int32")
      .Output("values: T")
      .Output("indices: int32")
      .Attr("sorted: bool = true")
      .Attr("T: realnumbertype")
      .SetShapeFn(TopKShapeFn);

  其映射的CANN算子TopK的定义如下所示：

  REG_OP(TopK)
      .INPUT(x, TensorType::RealNumberType())
      .INPUT(k, TensorType({DT_INT32}))
      .OUTPUT(values, TensorType::RealNumberType())
      .OUTPUT(indices, TensorType({DT_INT32}))
      .ATTR(sorted, Bool, true)
      .ATTR(largest, Bool, true)
      .ATTR(dim, Int, -1)
      .OP_END_FACTORY_REG(TopK)

  从以上定义中可以看出，CANN算子TopK的属性largest与dim在TensorFlow的TopKV2算子中未定义，所以我们需要在ParseParamsByOperatorFn的回调函数TopKMappingFn中实现对这两个属性的赋值，实现如下所示。

  Status TopKMappingFn(const ge::Operator &op_src, ge::Operator& op) {
    // 使用AutoMappingFn函数实现可以对应的属性的映射
    AutoMappingFn(op_src, op) != SUCCESS) {
      return FAILED;
    }
    // 为CANN算子TopK的dim属性赋初始值
    int32_t dim = -1;
    op.SetAttr("dim", dim);
    // 为CANN算子TopK的largest属性赋初始值
    bool largest = true;
    op.SetAttr("largest", largest);
    return SUCCESS;
  }
  REGISTER_CUSTOM_OP("TopK")
      .FrameworkType(TENSORFLOW)
      .OriginOpType("TopKV2")
      .ParseParamsByOperatorFn(TopKMappingFn)     //调用TopKMappingFn函数进行属性解析
      .ImplyType(ImplyType::TVM);

• 数据排布格式敏感算子
  • 对于在TensorFlow原图中有数据排布格式定义的算子，如果CANN算子规格中要求的张量数据排布格式定义与TensorFlow原图中张量的数据排布格式定义一致，则AutoMappingByOpFn函数会自动进行数据排布格式的处理。
  • 对于在使用TensorFlow原图中没有数据排布格式定义，或者数据排布格式与CANN中算子规格要求不一致的算子，则需要在回调函数ParseParamByOpFunc中强制设置张量的Format以及OriginalFormat（当前版本Format与OriginalFormat设置为一致即可）。
    例如针对Conv2D算子，TensorFlow接口规格以及CANN算子规格中的第二个输入filter支持的数据排布格式都为HWCN，但在TensorFlow原图（使用TensorFlow原生接口构造的图）中，输入filter的数据排布格式会直接继承属性“data_format”的值，所以需要在Parser函数中将filter的数据排布格式设置为期望的HWCN。如下所示：

    const int kInputFilter = 1
    Status ParseParamsConv2D(const ge::Operator &op_src, ge::Operator& op) {
        AutoMappingByOpFn(op_src, op);
        TensorDesc org_tensor_w = op.GetInputDesc(kInputFilter);
        org_tensor_w.SetOriginFormat(ge::FORMAT_HWCN);
        org_tensor_w.SetFormat(ge::FORMAT_HWCN);
        auto ret = op.UpdateInputDesc(kInputFilter, org_tensor_w);
        if (ret != ge::GRAPH_SUCCESS) {
             return FAILED;
        }
        return SUCCESS;
    }

    注意：以上仅为辅助描述此场景的示例，实际AI CPU的数据排布格式仅支持NHWC。

• 动态输入/输出算子
  对于存在动态输入/输出的算子，需要在插件的回调函数ParseParamByOpFunc中使用AutoMappingByOpFnDynamic实现TensorFlow算子和CANN算子的匹配。

  Status BoostedTreesBucketizeMapping(const ge::Operator& op_src, ge::Operator& op) {
    if (AutoMappingByOpFn(op_src, op) != SUCCESS) {
      return FAILED;
    }
  
    std::string attr_name = "num_features";
    std::vector<std::string> dynamic_inputs {"float_values", "bucket_boundaries"};
    std::string dynamic_output = "y";
  
    vector<DynamicInputOutputInfo> dynamic_name_attr_value;
  
    // input dynamic tensor
    for (std::string input_name : dynamic_inputs) {
        DynamicInputOutputInfo name_attr(kInput, input_name.c_str(), input_name.size(), 
                attr_name.c_str(), attr_name.size());
        dynamic_name_attr_value.push_back(name_attr);
    }
  
    // output dynamic tensor
    DynamicInputOutputInfo name_attr(kOutput, dynamic_output.c_str(), dynamic_output.size(), 
            attr_name.c_str(), attr_name.size());
    dynamic_name_attr_value.push_back(name_attr);
  
    AutoMappingByOpFnDynamic(op_src, op, dynamic_name_attr_value);
  
    return SUCCESS;
  }

多对一映射

前面的插件实现中，我们讲的将TensorFlow网络中算子映射为CANN算子的场景都是一对一映射的场景，在某些场景下，为了提高计算性能，往往需要将TensorFlow网络中的多个小算子融合并映射为CANN的一个大算子，以充分利用硬件资源进行加速。

这种场景下，我们首先需要设计并开发Scope融合规则，然后再实现融合算子适配插件，与普通的一对一映射的算子适配插件实现不同，融合算子适配插件的定义如下所示：

REGISTER_CUSTOM_OP("OpType")
    .FrameworkType(TENSORFLOW)              // 原始框架为Tensorflow
    .OriginOpType("OriginOpType")   // 算子在原始框架中的类型，和GenerateFusionResult的SetType的内容保持一致
    .FusionParseParamsFn(DecodeBboxV2ParseParams)  // 用来注册解析融合算子属性的函数
    .ImplyType(ImplyType::TVM);             // 指定算子的实现方式，ImplyType::TVM表示该算子是TBE算子

• REGISTER_CUSTOM_OP中注册的OpType为映射的CANN算子类型。
• OriginOpType中注册的为Scope融合规则开发时，设置的融合算子的结果类型。
• FusionParseParamsFn为注册解析融合算子属性的函数，接口定义可参见FusionParseParamsFn（Overload）。

此处不对Scope融合规则开发以及融合算子适配插件的开发做详细说明，详细的开发指导可参见《TensorFlow Parser Scope融合规则开发指南》。

多对多映射

若您需要将TensorFlow网络中的多个算子融合并映射为CANN的多个算子，此种场景下，您无需进行算子适配插件的开发，仅需实现Scope融合规则，并在Scope融合规则开发时，设置融合结果，将其映射为对应的CANN算子即可。详细的开发指导可参见《TensorFlow Parser Scope融合规则开发指南》。