Overview

Custom operators allow you to connect your own kernels (implemented in Ascend C, Triton, or other devices) into the Graph Engine (GE) compilation and execution pipeline. This enables you to leverage the performance benefits of graph offloading without modifying the GE framework codebase.

Application scenarios:

The frontend framework (PyTorch or TensorFlow) introduces operators that are not yet supported by GE and need to be executed on the NPU.
Custom kernel implementations (such as fusion operators or application-specific optimizations) need to be integrated into the GE compilation and offloading workflow.

Comparison with GE Ascend IR operators:

Dimension	Ascend IR Operators	Custom Operators
Registration method	Registered internally within GE via REG_OP and engine registration.	Registered via REG_OP + REG_AUTO_MAPPING_OP; delivered and loaded as a .so library.
Build artifact	Bundled with the GE release.	Packaged as a standalone .so file, loaded via ASCEND_CUSTOM_OPP_PATH.
Execution mode	Scheduled and executed directly by GE engines (for example, DNNEngine).	Executed via framework callbacks to the user-implemented Execute/Compile functions.
Update mode	Updated as part of the GE release cycle.	Decoupled from GE; the .so library can be replaced independently at any time.

Core Concepts

API System

All custom operator implementation classes must inherit from BaseCustomOp and flexibly combine the following capability APIs as required:

BaseCustomOp (common base class)
├── EagerExecuteOp     → Execute(ctx)          Runtime execution
├── ShapeInferOp      → InferShape (ctx)       Shape inference
│                        → InferDataType (ctx)    DataType inference
├── CompilableOp       → Compile (ctx)          Online build
├── PortableOp         → Serialize (buf)        Serialization to OM
│                        → Deserialize (buf)      Deserialization from OM
└── ArgsUpdater        → UpdateHostArgs (ctx)   I/O address update

Each API corresponds to an independent callback trigger. GE invokes the corresponding callback in specific phases of graph build or execution.

Registration mechanism

Custom operators require two registration steps:

REG_OP: defines the input, output, and attribute .proto file of an operator, used by the graph construction side to create nodes.

// Register the GE operator prototype, declaring the inputs, outputs, and supported data types of the operator.
// MyCustomOp is the operator type name, which must be the same as the class name in REG_AUTO_MAPPING_OP.
REG_OP(MyCustomOp)
    .INPUT (x, TensorType ({DT_FLOAT, DT_FLOAT16})) // Input x.
    .INPUT (y, TensorType ({DT_FLOAT, DT_FLOAT16})) // Input y.
    .OUTPUT(z, TensorType({DT_FLOAT, DT_FLOAT16})) // Output z.
    .OP_END_FACTORY_REG(MyCustomOp);

REG_AUTO_MAPPING_OP: maps the C++ implementation class to the operator type name. GE looks up this type name and instantiates the corresponding class:

// Custom operator implementation class, which inherits EagerExecuteOp to get the runtime execution capability.
class MyCustomOp : public EagerExecuteOp { /* ... */ };
// Register the implementation class with GE. When GE encounters a node of the MyCustomOp type in the graph, it instantiates this class and calls back its Execute method.
REG_AUTO_MAPPING_OP(MyCustomOp);

Naming constraints: The operator type name in REG_OP, the class name in REG_AUTO_MAPPING_OP, and the operator type used during graph construction must be the same.

Deliverables and OPP Package

The custom operator is delivered as a .so file loaded by GE. The path is specified via the ASCEND_CUSTOM_OPP_PATH environment variable. Recommended OPP package directory structure:

<custom_opp_root>/
├── op_graph/
│   ├── include/
│   │   └── my_custom_op.h          // REG_OP proto header file
│   └── lib/
│       └── <os>/<arch>/
│           └── libcust_opapi.so     // Deliverable (.so)
└── framework/                        // Framework deliverables (optional)
    └── tensorflow/
        └── npu_supported_ops.json

Configuration:

export ASCEND_CUSTOM_OPP_PATH="<custom_opp_root>:$ASCEND_CUSTOM_OPP_PATH"

Execution Scenarios

The development of custom operators requires different API combinations based on specific scenarios. This part describes the typical application scenarios and their key development key considerations. The following decision table describes the execution scenarios.

Scenario	Recommended API Combination	Typical Pipeline	Example
Online execution of dynamic graphs	EagerExecuteOp + ShapeInferOp (optional)	Graph construction → Direct execution	ascendc_add_custom
Online build + Online execution	EagerExecuteOp + CompilableOp + ShapeInferOp (optional)	Graph construction → Compilation → Execution	-
Offline OM model offloading	EagerExecuteOp + CompilableOp + ShapeInferOp + PortableOp	AIR → ATC → OM → ACL	compilable_add_custom
Data-dependent shape	EagerExecuteOp + ShapeInferOp	Graph construction → Execution + Shape buffer write-back	data_dependent_shape_custom

Scenario 1: Online execution of dynamic graphs
This is the simplest scenario. The kernel has been precompiled (for example, Ascend C .asc bundled with the library and Triton .npubin). You only need to load and launch the kernel in Execute.
```
Graph construction → GE callback to Execute → Obtain the input → Allocate the output → Launch the kernel → Return
```
Key points:
- Only EagerExecuteOp needs to be implemented.
- ShapeInferOp is optional. If the output shape is the same as the input shape and has already been declared in REG_OP, ShapeInferOp can be omitted.
- The kernel binary is compiled together with the .so file, or loaded from a file within Execute.
Scenario 2: Online compilation + online execution
The kernel source code needs to be compiled into a device binary through runtime compilation (RTC) in the GE compilation phase.
```
Graph construction → GE callback to Compile (RTC compiles the kernel)
     → GE callback to Execute (loading the binary and launching the kernel)
```
Key points:
- In the Compile callback, you can use ctx->GetInputTensor(i) to obtain the complete metadata of the inputs, such as the shape, data type, and format.
- It is recommended that the binary key be generated based on the shape to support multi-shape and multi-kernel management.
- The build artifact should be cached in a member variable of the implementation class for Execute to use.
Scenario 3: Offline OM model offloading
Based on scenario B, the build artifact needs to be serialized to an OM model file, and can later be deserialized when the OM model is loaded.
```
Graph construction → ATC offline build → Callback to Compile (RTC build)
                     → Callback to Serialize (binary data written to the OM)
     → ACL loads the OM model    → Callback to Deserialize (binary restored from OM)
                     → Callback to Execute (kernel launched)
```
Key points:
- The PortableOp must be implemented. The buffer formats of Serialize and Deserialize are user-defined. GE only passes the buffer through without parsing it.
- ShapeInferOp is called in the OM build phase to infer the output shape and data type.
- Serialization of multiple copies of binary files is supported (managed by key).
Scenario 4: Data-dependent shape (three types of operators)
The output shape depends on the runtime input data (such as Where and NonZero). Only an upper bound can be determined during the build.
```
InferShape → Output the upper bound of the shape (including unknown dimensions).
Execute → Allocate the output based on the upper bound and allocate the shape buffer.
           → Launch the kernel (The kernel writes the actual shape back to the shape buffer.)
           → Synchronize stream → D2H reads the shape buffer back → Update the output shape
```
Key points:
- In InferShape, use -1 to indicate an unknown dimension.
- In Execute, allocate the output based on the maximum shape and additionally allocate a shape buffer (using MallocWorkSpace).
- After the kernel is executed, synchronize the stream and read the shape buffer back.
- Shape buffer format: [ndim, d0, d1,...] (an array of uint64_t values).

Parent topic: Integrating Custom Operators into a Graph for PyTorch