Implementing Pass Using Python Based on Pattern Mapping

This section describes how to develop custom fusion passes in Python. Python fusion passes leverage the same GE matching and replacement mechanism as C++ fusion passes, but offer faster development cycles and easier runtime integration, significantly reducing development overhead. While the core mechanism is identical – both follow the same three-step "Match → Decide → Replace" workflow – Python implementations offer distinct advantages described below.

Dimension	Python Implementation	C++ Implementation
Development mode	Directly load .py files without compilation.	A .so dynamic library needs to be compiled.
Pattern definition	The @pattern decorator is supported, which can be directly described using expressions.	Use EsGraphBuilder to build a pattern graph.
Replacement definition	For simple replacements, you can simply return an expression, and that expression serves directly as the replacement.	You must manually create a replacement graph.
Iteration speed	Simply re-trigger execution or reload after modifying files.	Recompilation and reinstallation are required after modifications.
Integration method	Configure the ASCEND_GE_PY_PASS_PATH environment variable.	Place a compiled file in a specified directory.
Application scenario	Quick prototype development, rule verification, and iterative debugging	Production environments with high-performance requirements

The logic of fusion consists of three phases below:

Matching: A pattern defines a graph structure for subgraph matching and search.
Decision-making: The system determines whether fusion can be performed based on the matching results and specific conditions.
Replacement: If fusion is permitted, the system replaces the graph with the fused structure.

Figure 1 shows the logical architecture. You can implement custom fusion pass by inheriting the base classes provided by the GE and overriding their methods. The GE provides two types of base classes for inheritance in Python based on different application scenarios:

General subgraph fusion (one-to-one or complex topology replacement): A complete subgraph structure needs to be matched and replaced with another subgraph.
The PatternFusionPass class is inherited to implement the custom fusion pass, and the @register_fusion_pass decorator is used to register the pass to a specified phase.
Single-node replacement (one-to-N replacement): The DecomposePass class is inherited to implement the custom fusion pass, and the @register_decompose_pass decorator is used to register the pass to a specified phase.

The two scenarios use similar methods but differ in how they match and replace objects. The first scenario matches subgraphs and replaces complex topologies, while the latter matches nodes and replaces them with multiple nodes.

Scenarios

Fusion pass development in the general subgraph fusion scenario (one-to-one or complex topology replacement)

This section introduces the core data structure PatternFusionPass involved in this scenario, describes the core methods that need to be overridden during development, and explains how to register a pass and specify its execution phase. The PatternFusionPass class is declared as follows:

        
             class PatternFusionPass:
    def patterns(self) -> List[Pattern]:
        ...

    def meet_requirements(self, match_result: MatchResult) -> bool:
        return True

    def replacement(self, inputs, match_result: MatchResult = None):
        ...

You can develop a custom fusion pass by inheriting the PatternFusionPass class and overriding its core methods. The methods are described in the table below.

Method	Description	Must Override
patterns	Define what to match: Approach 1 (the @pattern decorator): Define the pattern topology to match in the target graph using Python expressions. The @pattern decorator is applicable to most common topologies. It automatically captures externally referenced inputs and returned pattern outputs, but does not automatically capture intermediate tensors not returned as outputs. Approach 2: Manually build a pattern graph and return a pattern list. If the intermediate tensor not returned as an output needs to be read by meet_requirements or replacement, you need to manually build a pattern and call the capture_tensor method to mark the intermediate tensor to be read. For simple rules like Add(x, 0) or MatMul + Add, prefer @pattern for shorter, more intuitive code.	Yes. Choose either of the two approaches.
meet_requirements	Check whether replacement is possible: This method filters the matched graph structures based on conditions after the pattern runs. The match result is input, and a Boolean value is returned.	No. By default, True is returned.
replacement	Define what to replace with: This method defines the replacement structure. The inputs object and match_result are passed as arguments, and the replacement structure is returned.	Yes

Method

Description

Must Override

patterns

Define what to match:

Approach 1 (the @pattern decorator): Define the pattern topology to match in the target graph using Python expressions.
The @pattern decorator is applicable to most common topologies. It automatically captures externally referenced inputs and returned pattern outputs, but does not automatically capture intermediate tensors not returned as outputs.
Approach 2: Manually build a pattern graph and return a pattern list.
If the intermediate tensor not returned as an output needs to be read by meet_requirements or replacement, you need to manually build a pattern and call the capture_tensor method to mark the intermediate tensor to be read.

For simple rules like Add(x, 0) or MatMul + Add, prefer @pattern for shorter, more intuitive code.

Yes. Choose either of the two approaches.

meet_requirements

Check whether replacement is possible: This method filters the matched graph structures based on conditions after the pattern runs. The match result is input, and a Boolean value is returned.

No. By default, True is returned.

replacement

Define what to replace with: This method defines the replacement structure. The inputs object and match_result are passed as arguments, and the replacement structure is returned.

Yes

Patterns

Using the @pattern decorator (recommended)

Using the @pattern decorator can greatly simplify the definition of patterns. You can directly use Python expressions to describe graph structures:

            
                 @register_fusion_pass(name="MyPass", stage=PassStage.BEFORE_INFER_SHAPE)
class MyPass(PatternFusionPass):
    @pattern
    def my_pattern(self, inputs):
        return inputs[0] + inputs[1]  # The expression is the pattern.

    def meet_requirements(self, match_result):
        # (Optional) Condition-based filtering
        return True

    def replacement(self, inputs):
        return inputs[0]  # The expression is the replacement.

Manually building a pattern (with the same method as that in C++)

For finer-grained control, you can build patterns manually, consistent with the C++ approach:

            
                 from ge.es.graph_builder import GraphBuilder
from ge.passes import create_pattern, create_replacement

def patterns(self):
    builder = GraphBuilder("pattern")
    a, b, c = builder.create_inputs(3)
    matmul = MatMul(a, b)
    add = matmul + c
    pat = create_pattern(builder.build_and_reset([add]))
    pat.capture_tensor(matmul)
    pat.capture_tensor(add)
    return [pat]

def replacement(self, match_result):
    builder = GraphBuilder("replacement")
    a, b, c = builder.create_inputs(3)
    gemm = GEMM(a, b, c, builder.create_scalar_float(1.0), builder.create_scalar_float(1.0))
    return create_replacement(builder.build_and_reset([gemm]))

meet_requirements

The matching results obtained by patterns are filtered in meet_requirements. Each matching result of the match_result type is used as an input parameter of meet_requirements. You can filter the information from match result. Returned Boolean values are used as the basis for determining whether to replace the matching results.

           
                def meet_requirements(self, match_result):
    # Traverse all nodes in matching results.
    for node in match_result.get_matched_nodes():
        # If the current node is a constant node        
        if node.type == "Const":
            return _is_zero(node.get_attr("value"))
    return False

match_result is a matching result class, which contains information such as matched nodes and edges. You can use the member functions below to filter matches.

Method	Description
get_matched_nodes()	Obtains the list of all matched nodes.
get_captured_tensor(capture_index)	Obtains the captured tensor of a specified index.

The following is an example of using the get_captured_tensor member function to check whether the ReLU output is a dynamic shape:

          
               def meet_requirements(self, match_result):
    from ge.graph.types import TensorDesc

// Try to obtain the first captured output tensor from match_result.
    try:
        relu_output = match_result.get_captured_tensor(0)
        # Obtain the description of the output tensor.
        relu_out_tensor_desc = relu_output.node.get_output_desc(relu_output.index)
        # Check whether the shape is dynamic.
        if relu_out_tensor_desc.shape.size != -1:
            return False
        return True
    except Exception:
        return False

replacement

The replacement method defines the target structure to replace the subgraph matched by patterns where meet_requirements returns True. Python provides two approaches to define a replacement:

Approach 1: Returning an expression directly (simple scenario)

For simple replacements, you can directly return an expression. GE automatically builds a replacement graph.

            
                 def replacement(self, inputs):
    # Return the 0th input directly, which is equivalent to deleting a node.
    return inputs[0]

Alternatively, use an operator expression:

            
                 def replacement(self, inputs):
    a, b, c = inputs[:3]
    # Return a GEMM expression.
    return GEMM(a, b, c, 1.0, 1.0)

Approach 2: Complex scenario (accepting both inputs and match_result)

If you want to use an expression to simplify the replacement construction and read the attributes of matched nodes, you can accept both parameters:

            
                 def replacement(self, inputs, match_result):
    a, b, c = inputs[:3]
    transpose_a = False
    transpose_b = False
    for node in match_result.get_matched_nodes():
        if node.type not in ("MatMul", "BatchMatMulV2"):
            continue
        try:
            transpose_a = bool(node.get_attr("transpose_x1"))
            transpose_b = bool(node.get_attr("transpose_x2"))
        except RuntimeError:
            pass
        break
    return GEMM(a, b, c, 1.0, 1.0, transpose_a, transpose_b)

Registration of a custom fusion pass

Use the @register_fusion_pass decorator to register a class with GE.

        
             @register_fusion_pass(name="MyPass", stage=PassStage.BEFORE_INFER_SHAPE)
class MyPass(PatternFusionPass):
    ...

The name must be unique. stage indicates the execution phase (see PassStage). For initial development, PassStage.BEFORE_INFER_SHAPE is recommended, as the replacement can still go through GE's subsequent unified shape inference process.

Fusion pass development (one-to-N node replacement)

In this scenario, the base class inherited by the pass is DecomposePass. Since the replaced structure is a single node, patterns do not need to be manually defined. Instead, the operator type is specified via the op_types parameter during registration:

        
             from ge.passes import DecomposePass, PassStage, register_decompose_pass

@register_decompose_pass(
    name="PythonMyDecomposePass",
    stage=PassStage.AFTER_INFER_SHAPE,
    op_types=["Conv2D"],
)
class PythonMyDecomposePass(DecomposePass):
    def meet_requirements(self, node):
        return node.get_attr("groups") != 1

    def replacement(self, node):
        # Return the replacement graph consisting of basic operators.
        ...

op_types determines which node types GE will route to this pass, and meet_requirements further filters which specific nodes should actually be replaced.

Registration of the custom fusion pass

The following is an example of using the @register_decompose_pass decorator (see register_decompose_pass) to initialize CustomOne2NPass with Conv2D as op_types. In the example, the execution phase for the registered pass is kAfterInferShape:

        
             @register_decompose_pass(name="CustomOne2NPass", stage=PassStage.AFTER_INFER_SHAPE, op_types=["Conv2D"])
class CustomOne2NPass(DecomposePass):
    ...

Examples

The following uses the Add(x, 0) structure as an example (using the @pattern decorator) to describe how to modify a graph using a custom fusion pass in Python. For the complete sample code, see AddZeroPass Python Sample. The source repository also provides the following examples for reference in various scenarios:

The graph structures before and after the modification are as follows:

     
          x ----\
       Add ---- out      ==>      x ---- out
0 ----/

Import necessary modules.

       
            from math import fabs

from ge.graph.types import DataType
from ge.passes import PassStage, PatternFusionPass, pattern, register_fusion_pass

Implement auxiliary functions.

       
            def _scalar_value(value):
    while isinstance(value, list):
        if len(value) != 1:
            return None
        value = value[0]
    return value

def _is_zero(tensor):
    value = _scalar_value(tensor.data)
    if value is None:
        return False
    if tensor.data_type == DataType.DT_FLOAT:
        return fabs(float(value)) < 1e-6
    if tensor.data_type == DataType.DT_DOUBLE:
        return fabs(float(value)) < 1e-15
    if tensor.data_type == DataType.DT_INT32:
        return int(value) == 0
    return False

Implement a custom pass.

       
            @register_fusion_pass(name="PythonAddZeroPass", stage=PassStage.BEFORE_INFER_SHAPE)  ## Register a custom fusion pass and specify the phase to be called.

class PythonAddZeroPass(PatternFusionPass):
    @pattern
    def add_zero(self, inputs):
        return inputs[0] + 0  # Define the structure to be matched: the 0th input plus 0.

    def meet_requirements(self, match_result):
        # Check whether the matched constant node is 0.
        for node in match_result.get_matched_nodes():
            if node.type != "Const":
                continue
            return _is_zero(node.get_attr("value"))
        return False

    def replacement(self, inputs):
        return inputs[0] # Return the 0th input directly, which is equivalent to deleting Add.

@pattern decorator: The inputs[0]+0 expression describes the structure to be matched: the 0th external input plus a constant 0. GE automatically identifies this expression and builds the corresponding pattern.
meet_requirements method: Traverse the nodes in the matching result, find the Const node, and check whether its value is 0. Only the matches that meet the conditions will be replaced.
replacement method: The method simply returns inputs[0], indicating that the entire Add node is replaced by the 0th input, effectively removing the redundant Add(x, 0) operation.

How to Use a Custom Pass

This part describes how to load the Python pass to the GE compilation phase after the custom pass is implemented.

Set CANN environment variables.
```
source /usr/local/Ascend/cann/set_env.sh
```
/usr/local/Ascend/ is the default installation path for the root user. Replace it with the actual one.
Set the environment variable for loading the pass.
Specify the path from which GE loads the Python pass.
```
export ASCEND_GE_PY_PASS_PATH=/path/to/my_pass.py
```
It can also point to a directory.
```
export ASCEND_GE_PY_PASS_PATH=/path/to/pass_dir/
```
Use colons (:) to separate multiple paths.
```
export ASCEND_GE_PY_PASS_PATH=/path/to/a.py:/path/to/pass_dir/
```
For details about ASCEND_GE_PY_PASS_PATH, see Environment Variables.
Perform compilation. The supported entries for building model files include, but are not limited to:
To check whether the custom pass takes effect, dump the graph before model build by setting the DUMP_GE_GRAPH environment variable and build the model from the following entries:
- Offline compilation: In offline scenarios, you are advised to use pyatc to trigger compilation. The command line parameters of pyatc and atc are the same, but pyatc runs in the current Python interpreter process to facilitate Python pass loading. The command is as follows:
```
pyatc --model=./model.onnx --framework=5 --soc_version=xxx --output=./model
```
  For details about how to query the value of --soc_version, see 2.
  
  For details about how to use pyatc, see Pyatc APIs.
- Online compilation: Use the torch_forward.py script to perform online inference. For details, see Sample Program Execution.

Result Verification

Enable graph dump.
It is recommended that graph dump be enabled during each development.
```
export DUMP_GE_GRAPH=1
```

Check the graphs before and after the fusion.
After the dump environment variable is set, the graph files such as ge_onnx*.pbtxt are generated in the current path after the program is executed. You can obtain the following two graphs and use visualization software such as Netron to view the graphs. The pass execution phase is executed before InferShape in this case.
- ge_onnx_xxxx_PreRunBegin.pbtxt: graph before fusion
- ge_onnx_xxxx_RunCustomPassBeforeInfershape.pbtxt: graph after fusion
After the execution is successful, information similar to the following is displayed in the log (the name of input_0 may be different):
```
[PythonAddZeroConstValueMatchPass] matched=PythonAddZeroConstValueMatchPass_add_zero_pattern captured=input_0:0
[PythonAddZeroPass] matched=PythonAddZeroPass_add_zero_pattern captured=input_0:0 const_dtype=DT_FLOAT zero=True
```

If the pass is not matched, perform the checks below in sequence.

Symptom	Possible Cause	Check Method
The Python file is not loaded.	ASCEND_GE_PY_PASS_PATH is not set, the path does not exist, or the file name extension is not .py.	Check the environment variables and paths.
The class has been loaded, but the pass is not executed.	The registration decorator is not used, or the registration phase is incorrect.	Check the @register_fusion_pass or @register_decompose_pass decorator.
The pattern is not matched.	The operator type, number of inputs, or output boundary is inconsistent.	Compare against the actual topology in the dumped graph.
The pattern is matched, but replacement is skipped.	meet_requirements returns False.	Print the attributes of the matched node.
The graph is abnormal after replacement.	The replacement output does not cover the tensors required by external consumers.	Check the boundary rules based on Fusion Pattern Pass Mechanism.

Parent topic: Modifying a Graph Using a Custom Pass