Plugin Development (ONNX)

Overview

This section describes how to develop an operator plugin capable of mapping an operator trained in a third-party framework to one adapted to Ascend AI Processor, and how to register the operator information with GE. During the execution of ONNX-based network, the plugin information in GE is loaded to parse and map the operators in the network to those that can adapt to Ascend AI Processor.

The operators adapted to Ascend AI Processor are referred to as CANN operators.

Principles

The implementation workflow of an operator plugin includes registering the operator types of CANN operators and the operators in the original framework, and mapping the operator attributes in the original framework to attributes of CANN operators. Operator mapping is implemented by the Parser module. Figure 1 shows the plugin implementation workflow in the network execution scenario.

Figure 1 Process for implementing an operator plugin

1. GE receives a graph, which represents the topology of an original network model trained under a third-party framework. Then, GE is initialized.
2. GE loads the .so file of the operator plugin from the Register module, which is stored in the opp/built-in/framework/ directory in the CANN component directory.
3. The Register module reads operator information in the operator plugin .so file and registers the information with the map file of the operator plugin. (Information about all operator plugins is stored in one map file.)
4. GE requests the Parser module to call the Parser method.
5. The Parser module obtains the corresponding Parser function from the map file of the operator plugin based on the operator type (OpType), and returns the implementation function ParseParamsByOperatorFn to the Parser module. The Parser module maps the attributes of a third-party operator to attributes of a CANN operator (that is, the attributes defined in the operator prototype). In this way, a third-party operator is mapped to a CANN operator.
6. Subsequent operations are performed, including graph preparation, partition, and optimization. Then a network model that adapts to Ascend AI Processor is generated.

Plugin Implementation

GE provides the REGISTER_CUSTOM_OP macro to register an operator based on the specified operator name.

The following code shows how to customize an ONNX operator.

#include "register/register.h"
#include "graph/operator.h"
#include "json.hpp"
namespace domi
{
REGISTER_CUSTOM_OP("OpType")
    .FrameworkType(ONNX) 
    .OriginOpType("OriginOpType")
    .ParseParamsByOperatorFn(ParseParamByOpFunc)   // Registers a function for parsing operator attributes.
    .ImplyType(ImplyType::TVM);  // TBE operator: ImplyType::TVM; AI CPU operator: ImplyType::AI_CPU
}

Mapping Operators to Subgraphs (One-to-Many)

The following scenarios may occur during adaptation development:

• There's no matching CANN operator for a certain ONNX operator, but you can achieve its function by combining multiple CANN operators.
• The implementation of ONNX operators is different from that of CANN operators. For example, some attribute of an ONNX operator is a constant input of the corresponding CANN operator.
  For example, the prototype of the ArgMax operator in ONNX is as follows.

  Table 1

  ONNX ArgMax
  Attributes	axis: The axis along which the arg indices are computed.
  	keepdims: Whether to keep the reduced dimension. Defaults to 1, indicating that the reduced dimension is kept.
  	select_last_index: Whether to select the last index or the first index if {name} appears in multiple indices. Defaults to False (first index).
  Inputs	data: An input tensor.
  Outputs	reduced: Reduced output tensor with the integer data type.

  The prototype of the mapped CANN operator ArgMaxV2 is as follows.

  Table 2

  CANN ArgMaxV2
  Inputs	x: A multi-dimensional Tensor of type float16, float32, or int16.
  	dimension: A Scalar of type int32, specifying the index with the largest value.
  Attributes	dtype: The output type, either "int32" or "int64". Defaults to "int64".
  Outputs	y: A multi-dimensional Tensor of type int32 or int64, specifying the index with the largest value. The dimension is one less than that of "x".

  According to the comparison between the prototypes of the two operators, the axis attribute of the ONNX operator ArgMax is defined as a constant input of the CANN operator ArgMaxV2.

In these scenarios, one operator in the original ONNX framework needs to be mapped to multiple operators in the CANN. During plugin implementation, you need to build the mapped CANN operators to a subgraph, and then map the ONNX operator to the subgraph. You can obtain implementation samples from the cplusplus\level1_single_api\4_op_dev\1_custom_op\framework\onnx_plugin directory in the open-source samples repository to understand the implementation method and verify the mapping result. The following describes how to convert the ONNX AddN operator into two CANN Add operators.

The operator registration code is as follows:

REGISTER_CUSTOM_OP("PartitionedCall")
    .FrameworkType(ONNX)
    .OriginOpType("ai.onnx::11::AddN")
    .ParseParamsByOperatorFn(ParseParamsAddn)
    .ParseOpToGraphFn(ParseOpToGraphAddn)
    .ImplyType(ImplyType::TVM);

The following describes only the differences from the common registration function:

• REGISTER_CUSTOM_OP("PartitionedCall"): registers the subgraph PartitionedCall, which is a general name of combined operators.

  If an operator is mapped to a subgraph, the value is fixed at PartitionedCall.

• ParseParamsByOperatorFn(ParseParamsAddn): registers the ParseParamsAddn function for parsing custom operator parameters. In ParseParamsAddn, the parameters of the original ONNX operator are mapped to those of the PartitionedCall operator. You need to set the number of inputs and outputs of the PartitionCall node and set the original_type attribute to the operator type in the original framework.
• ParseOpToGraphFn(ParseOpToGraphAddn): registers the ParseOpToGraphAddn function for implementing the one-to-many subgraph mapping of the operator. In ParseOpToGraphAddn, the PartitionedCall operator is mapped to a subgraph, which is constructed according to Ascend Graph construction. For details about the ParseOpToGraphFn API, see ParseOpToGraphFn. For details about Ascend Graph construction, see Ascend Graph Developer Guide.

The following is an example of implementing the callback function for mapping the AddN operator parameters to the PartitionedCall operator parameters:

// Example of the ParseParamsByOperatorFn callback function:
Status ParseParamsAddn(const ge::Operator&op_src, ge::Operator&op_dest) {
  // 1. Align the number of inputs and outputs of the PartitionCall node (op_dest) with that of the original node (op_src).
  ge::Operator op_ori = const_cast<ge::Operator&>(op_src);
  std::string in_name = "args";
  std::string in_value = "in_num";
  std::string out_name = "output";
  std::string out_value = "out_num";
  op_ori.SetAttr(in_value, 3);
  op_ori.SetAttr(out_value, 1);
  DynamicInputOutputInfo in_values(kInput, in_name.c_str(), in_name.size(), in_value.c_str(), in_value.size());
  DynamicInputOutputInfo out_values(kOutput, out_name.c_str(), out_name.size(), out_value.c_str(), out_value.size());
  AutoMappingByOpFnDynamic(op_ori, op_dest, {in_values, out_values});
  // 2. Inherit attributes, if any, from op_src to op_dest.
  ...
  // 3. Set the original_type attribute to the operator type in the original framework.
  op_dest.SetAttr("original_type", "ai.onnx::11::AddN");
  return SUCCESS;
}

// Example of the ParseOpToGraphFn callback function:
static ParseOpToGraphAddn(const ge::Operator&op, ge::Graph&graph) {
  // The index attribute of the Data node indicates the indexth input of the original node (op).
  auto data_0 = ge::op::Data().set_attr_index(0);
  auto data_1 = ge::op::Data().set_attr_index(1);
  auto data_2 = ge::op::Data().set_attr_index(2);
  // Create an add0 operator instance and set the operator inputs to data_0 and data_1.
  auto add0 = ge::op::Add("add0")
      .set_input_x1(data_0)
      .set_input_x2(data_1);
  // Create an add1 operator instance and set the operator inputs to data_2 and add0.
  auto add1 = ge::op::Add("add1")
      .set_input_x1(data_2)
      .set_input_x2(add0);

  // Set the inputs and outputs of the graph.
  std::vector<ge::Operator> inputs{data_0, data_1, data_2};
  // The output setting must be the same as that of the original node (op).
  std::vector<std::pair<ge::Operator, std::vector<size_t>>> output_indexs;
  output_indexs.emplace_back(add1, vector<std::size_t>{0});
  graph.SetInputs(inputs).SetOutputs(output_indexs);
  return SUCCESS;
}

Verify the one-to-many mapping result by referring to "Verifying the Mapping Between Operators and Subgraphs (One-to-Many)" in the README file under the cplusplus\level1_single_api\4_op_dev\1_custom_op directory in the open-source samples repository. The verification result is as follows:

In the following figure, the original ONNX model containing the AddN operator is displayed on the left, and the converted subgraph structure containing two CANN Add operators is displayed on the right.