Integrating Custom Operators into a Graph for ONNX

The ONNX frontend is different from PyTorch. If the ONNX Parser of GE cannot identify an operator type, it directly reports an error (PARAM_INVALID). Therefore, you need to compile an additional ONNX parsing plugin to map the attributes of the ONNX node to the attributes of the GE operator. The following figure shows the deliverables.

1. Phase 1: Operator definition

   In the operator definition phase, operator metadata, such as the inputs, outputs, and attributes, is declared for use in subsequent phases.

   Implementation by frontend framework:

   Frontend Framework	Definition Method	Deliverable	Description
   ONNX	NodeProto	Embedded in the ONNX model	Operator definition is natively included within the model itself.

2. Phase 2: Frontend integration

   This phase maps frontend operator types to GE operator types and converts attributes.

   Integration by frontend framework:

   Frontend Framework	Integration Method	Deliverable Involved	Key Point
   ONNX	Parser plug-in	libonnx_parser_plugin.so	ParseParamsFn must be manually written.

   ONNX operator type identifier format: "domain::version::OpType" (for example, "ai.onnx::11::Conv" or "com.example::1::MyOp"). The Parser searches for the mapping in OpRegistry based on this format. If the mapping cannot be found, the parsing fails.

   The ONNX parser plug-in requires you to:

   1. Use REGISTER_CUSTOM_OP to register the mapping from the ONNX operator type to the GE operator type.
   2. Implement ParseParamsFn to extract attributes from the ONNX NodeProto and set them for the GE operator.

   Parser plug-in code:

   // framework/onnx/add_custom_parser.cpp
   #include "register/register.h"
   #include "proto/onnx/ge_onnx.pb.h"
   
   // Attribute parsing function: Map the ONNX NodeProto attributes to the GE operator attributes.
   // op_src: proto message body on the ONNX node
   // op_dest: GE operator object, which is used to set attributes
   static domi::Status ParseMyOp(const google::protobuf::Message *op_src, ge::Operator &op_dest) {
     // Convert the message to the ONNX NodeProto type.
     auto *node = dynamic_cast<const ge::onnx::NodeProto *>(op_src);
     // Traverse all attributes of the ONNX node.
     for (const auto &attr : node->attribute()) {
       // Set the alpha attribute in the ONNX model to the GE operator.
       if (attr.name() == "alpha") {
         op_dest.SetAttr("alpha", attr.f());
       }
     }
     return domi::SUCCESS;
   }
   
   // Register the ONNX parser plug-in: Map the ONNX operator type to the GE operator type.
   // When the ONNX Parser encounters an unrecognized operator, it searches for the parser plug-in based on this registration information.
   REGISTER_CUSTOM_OP("MyCustomOp")
       .FrameworkType(domi::ONNX)                                    // Declare it as an ONNX framework plug-in.
       .OriginOpType({"ai.onnx::11::MyOp", "ai.onnx::13::MyOp"})     // ONNX operator type (multiple versions supported)
       .ParseParamsFn(ParseMyOp);                                    // Attribute parsing callback function
   

   Key methods of REGISTER_CUSTOM_OP

   Method	Function
   .FrameworkType(domi::ONNX)	Declares the operator as an ONNX framework plug-in
   .OriginOpType("domain::ver::Type")	Specifies the ONNX operator type (supporting a list of types, covering multiple versions).
   .ParseParamsFn(fn)	Attribute parsing function (Message* → Operator&)
   .ParseParamsByOperatorFn(fn)	Attribute parsing for automatic mapping paths (Operator → Operator)
   .ParseOpToGraphFn(fn)	Expands a single operator into a subgraph (optional)
   .ParseSubgraphPostFn(fn)	Post-processing for operators that contain subgraphs (optional)

   In ONNX scenarios, two .so files listed in the table below are required.

   Deliverable	Function	Directory
   ONNX parser plug-in (.so)	REGISTER_CUSTOM_OP + ParseParamsFn, responsible for attribute mapping from ONNX to GE	Plug-in directory under ASCEND_CUSTOM_OPP_PATH
   GE custom operator (.so)	REG_OP + REG_AUTO_MAPPING_OP + execution logic	op_graph/lib/<os>/<arch>/

   Deliverable structure example:

   custom_opp/
   ├── op_graph/
   │   ├── include/add_custom.h    # REG_OP
   │   └── lib/linux/x86_64/libcust_opapi.so  # GE deliverable
   ├── framework/
   │   └── onnx/
   │       └── libonnx_parser_plugin.so       # Parser plug-in (an independent .so file)

   Precautions:

   • The ONNX parser plug-in (.so) must be loaded through ASCEND_CUSTOM_OPP_PATH before ATC compilation. Otherwise, the parser cannot identify the custom operator type.
   • If the domain in OriginOpType is left empty, ai.onnx is used by default. A custom domain must be explicitly specified.
   • If an ONNX operator can be mapped to a combination of existing GE operators, you can use ParseOpToGraphFn to expand it into a subgraph. In this case, you do not need to compile the custom operator (.so) on the GE.
3. Phase 3: Graph construction and integration

   In the graph construction and integration phase, the frontend computational graph is converted into a GE computational graph.

   Graph integration path by frontend framework:

   Frontend Framework	Graph Integration Path	Conversion Mechanism
   ONNX	ONNX Model → GE Graph	Parser

   This phase is automatically completed by the framework, and you do not need to provide any additional deliverables. After the graph integration is complete, the GE graph already contains the custom operator node.

4. Phase 4: Graph compilation

   The callback APIs for graph compilation are triggered during GE compilation and are used for shape inference, kernel compilation, and serialization. Callback APIs and triggering timing are described in the table below.

   Callback API	Triggering Timing	Application Scenario
   InferShape	Compile-time shape inference	All scenarios (optional)
   InferDataType	Compile-time data type inference	All scenarios (optional)
   Compile	Compile-time RTC kernel compilation	Online compilation/Offline OM
   Serialize	Writing data to the OM during ATC compilation	Offline OM

   For details about the callback API examples, see Callback APIs.

5. Phase 5: Graph execution

   This phase is the core of custom operator integration. GE calls back the Execute API during execution.

   Callback APIs and triggering timing are described in the table below.

   Callback API	Triggering Timing	Application Scenario
   Execute	Runtime kernel execution	All scenarios (mandatory)
   Deserialize	Restoring binary during OM loading	Offline OM
   UpdateHostArgs	Updates upon I/O address changes	Static graph offloading

   For details about the callback API examples, see Callback APIs.

Mapping between ONNX graph integration phases and deliverables

Additional requirements:

• op_graph/include/*.h (REG_OP) is used by the parser to create nodes.
• The parser plug-in and GE deliverables must be loaded at the same time.