Integrating custom operators into a graph

Custom Operator Development for PyTorch

The following figure shows the process of developing custom Ascend C operators, adapting them to PyTorch, and integrate them into a graph.

1. Custom operator development
   Before integrating custom operators into a graph, you need to develop custom operators based on Ascend C.

   1. Create an operator project.
   2. Implement the operators, including operator prototype definition, operator implementation on the kernel, and tiling implementation on the host.
   3. Implement functions such as shape derivation and data type derivation to import operators to the GE graph.
   4. Build and deploy the operator project. Build and generate the custom operator installation package, install the operator package, and deploy the custom operators to the operator acceleration library.

   For details about the preceding process, see Integrating Operators into a GE Graph in Ascend C Operator Development Guide .

2. Adaptation to PyTorch
   After a custom operator is developed, you need to perform adaptation on the PyTorch side to import the custom operator to the graph. The operations are as follows:

   1. Register the operator with the PyTorch framework using OpPlugin.
   2. Register the meta function with torch_npu to support shape and type inference for operator integration into a graph.
   3. Register and implement the custom operator converter through TorchAir.

   For details about how to import a custom operator into a PyTorch graph, see "Importing a Custom Operator into a Graph Using a Plugin" in PyTorch Graph Mode Instructions (TorchAir) .

Custom Operator Development for TensorFlow

The following figure shows the process of developing a custom Ascend C operator, adapting the TensorFlow framework, and executing the operator into the graph.

1. Custom operator development
   Before integrating custom operators into a graph, you need to develop custom operators based on Ascend C.

   1. Create an operator project.
   2. Implement the operators, including operator prototype definition, operator implementation on the kernel, and tiling implementation on the host.
   3. Implement functions such as shape derivation and data type derivation to import operators to the GE graph.

   For details about the preceding process, see Integrating Operators into a GE Graph in Ascend C Operator Development Guide .

2. Develop the TensorFlow adaptation plugin. Register a custom operator and map the TensorFlow operator to the CANN operator by using REGISTER_CUSTOM_OP. In the scenario where a TensorFlow custom operator is mapped to a CANN operator, you also need to develop the TensorFlow custom operator. For details, see " AI Framework Operator Adaptation > TensorFlow Framework" in " Ascend C Operator Development Guide ."
3. Build and deploy the operator project. Build and generate the custom operator installation package, install the operator package, and deploy the custom operators to the operator acceleration library.
4. Call operators in the TensorFlow framework.

Constructing an Ascend graph

The following figure shows the development process of developing custom Ascend C operators and integrating them into an Ascend graph.

1. Custom operator development
   Before integrating custom operators into a graph, you need to develop custom operators based on Ascend C.

   1. Create an operator project and perform operator prototype definition, operator implementation on the kernel, and tiling implementation on the host.
   2. Implement the shape derivation and data type derivation functions to support the operator in graph mode.
   3. Build and deploy the operator project. Build and generate the custom operator installation package, install the operator package, and deploy the custom operators to the operator acceleration library.

   For details about the preceding process, see Integrating Operators into a GE Graph in Ascend C Operator Development Guide .

2. Ascend graph construction

   The method of constructing an Ascend graph based on a custom operator is the same as that of constructing an Ascend graph based on other built-in CANN operators. You can call the operator prototype. For details about the typical process, see Constructing an Ascend graph.