Integrating Custom Operators into a Graph for TensorFlow

The primary difference between TensorFlow and PyTorch scenarios is that the Adapter can automatically generate GE operator prototypes from the TF op_def, eliminating the need for developers to manually write REG_OP. The following figure shows the deliverables. For details about the TensorFlow scenario, see Integrating a Triton Operator into a Graph or triton_add_custom.

Each phase and its deliverables are described in detail as follows:

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.

   Frontend Framework	Definition Method	Deliverable	Description
   TensorFlow	TF op_def	libcustom_ops.so	Operator declaration in TensorFlow. REG_OP is automatically generated by GE.

2. Phase 2: Frontend integration

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

   Frontend Framework	Integration Method	Deliverable Involved	Key Point
   TensorFlow	Automatic mapping by the Adapter	npu_supported_ops.json	JSON configuration for operator visibility

   The following are required in TensorFlow scenarios:

   1. Custom operator .so on TensorFlow (to declare operator visibility).
   2. Deliverable .so on GE (only REG_AUTO_MAPPING_OP is required, and REG_OP is not required).
   3. npu_supported_ops.json (In TensorFlow scenarios, automatic mapping is performed through the Adapter. This deliverable is a key deliverable.)

   Differences from the PyTorch scenario: In TensorFlow scenarios, REG_AUTO_MAPPING_OP can automatically generate GE operator prototypes (including inputs, outputs, and attributes) from TensorFlow operator prototypes, and developers do not need to write REG_OP definitions. This further reduces the overhead for integrating custom operators for TensorFlow.

   Example of registering an operator on TensorFlow (using TensorFlow 1.15 as an example):

   #include "tensorflow/core/framework/op.h"
   #include "tensorflow/core/framework/op_kernel.h"
   #include "tensorflow/core/framework/shape_inference.h"
   #include "tensorflow/core/framework/common_shape_fns.h"
   
   using namespace tensorflow;
   
   // Register the operator prototype by using the REGISTER_OP API provided by TensorFlow.
   // After the registration, the GE Adapter can automatically generate the GE operator prototype from this definition. You do not need to manually write REG_OP.
   REGISTER_OP("AddCustom")
       .Input("x: T")
       .Input("y: T")
       .Output("z: T")
       .Attr("T: {float, double, int32}")
       .SetShapeFn(shape_inference::BroadcastBinaryOpShapeFn);
   
   // CPU implementation of the TensorFlow custom operator (CPUs are not supported and only declared.)
   // The actual computation is executed on NPUs through the GE graph.
   class AddCustomOp : public OpKernel {
   public:
       explicit AddCustomOp(OpKernelConstruction* context) : OpKernel(context) {}
       void Compute(OpKernelContext* context) override {
           OP_REQUIRES(context, false, errors::Unimplemented("AddCustomOp is not supported on CPU"));
       }
   };
   
   // Register the CPU kernel implementation. (This step is mandatory. Otherwise, TensorFlow cannot load the operator.)
   REGISTER_KERNEL_BUILDER(Name("AddCustom").Device(DEVICE_CPU), AddCustomOp);

   Example of npu_supported_ops.json:

   {
     "custom_ops": [
       {
         "op_name": "AddCustom",
         "op_type": "AddCustom"
       }
     ]
   }

   Deliverable structure example:

   custom_opp/
   ├── op_graph/
   │   └── lib/linux/x86_64/libcust_opapi.so  # GE deliverable (only REG_AUTO_MAPPING_OP is required)
   ├── framework/
   │   └── tensorflow/
   │       ├── libcustom_ops.so               # Operator declaration in TensorFlow
   │       └── npu_supported_ops.json         # Adapter configuration

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.

   Frontend Framework	Graph Integration Path	Conversion Mechanism
   TensorFlow	TensorFlow Graph → GE Graph	Adapter

   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 TensorFlow graph integration phases and deliverables