Python Fusion Pass Development Guide
【免费下载链接】geGE(Graph Engine)是面向昇腾的图编译器和执行器,提供了计算图优化、多流并行、内存复用和模型下沉等技术手段,加速模型执行效率,减少模型内存占用。 GE 提供对 PyTorch、TensorFlow 前端的友好接入能力,并同时支持 onnx、pb 等主流模型格式的解析与编译。项目地址: https://gitcode.com/cann/ge
This guide is for developers who want to write GE fusion passes in Python. It is recommended to first read the language-independent mechanism description: Fusion Pattern Pass Mechanism.
If you already understand the main workflow of "define pattern, match, filter, replacement, reconnect", you can start coding directly according to this guide.
1. Why Consider Python First
Python pass and C++ pass use the same GE matching and replacement mechanism, but Python is better suited for rapid development and runtime integration:
- Easy Integration: Configure
.pyfiles or directories toASCEND_GE_PY_PASS_PATH, GE will load at runtime during compilation phase, without compiling pass into.so. - Shorter Expression:
@patterncan use Python expressions to describe patterns, e.g.,return inputs[0] + 0. - Intuitive Replacement: Simple replacements can directly write
return inputs[0], without manually creating replacement graph. - Easy Iteration: Modify Python file and re-trigger compilation to validate, suitable for iterating rules first.
2. Minimal Example: Delete Add(x, 0)
Goal: ReplaceAdd(x, 0)in the graph withx.
x ----\ Add ---- out ==> x ---- out 0 ----/Python implementation:
from math import fabs from ge.graph.types import DataType from ge.passes import PassStage, PatternFusionPass, pattern, register_fusion_pass 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 @register_fusion_pass(name="PythonAddZeroPass", stage=PassStage.BEFORE_INFER_SHAPE) class PythonAddZeroPass(PatternFusionPass): @pattern def add_zero(self, inputs): return inputs[0] + 0 def meet_requirements(self, match_result): 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]This code does three things:
@patternmethod describes the structure to find: the 0th external input plus a constant.meet_requirementschecks if the matched constant is really 0.replacementreturns the 0th external input, equivalent to deleting the matchedAdd.
A complete runnable example is available at AddZeroPass Python Example.
3. Steps to Write a PatternFusionPass
3.1 Import Interfaces
Common imports:
from ge.passes import ( PassStage, PatternFusionPass, pattern, register_fusion_pass, )If writing aDecomposePass, also need:
from ge.passes import DecomposePass, register_decompose_passComplete interface documentation is at Python Passes API.
3.2 Register Pass
Use@register_fusion_passto register the class to GE:
@register_fusion_pass(name="MyPass", stage=PassStage.BEFORE_INFER_SHAPE) class MyPass(PatternFusionPass): ...namemust be unique.stageindicates execution stage. For initial development, usePassStage.BEFORE_INFER_SHAPE, because replacement can still go through GE's subsequent unified shape inference process.
3.3 Use @pattern to Define Structure to Match
@patternmethod receives aninputsobject. It represents the external input set of the pattern.
@pattern def add_zero(self, inputs): return inputs[0] + 0Hereinputs[0]is the 0th external input placeholder, not a fixed real node. When matching, GE will map the real tensor connected to this structure to it.
Multi-input scenarios can be written like this:
@pattern def matmul_add(self, inputs): a, b, c = inputs[:3] return MatMul(a, b) + cNotes:
inputs[i]will create theith input as needed.inputs[:N]is used to explicitly declare multiple consecutive inputs.@patternwill automatically capture visited external inputs and returned pattern outputs. Capture order is fixed: first capture external inputs by input index, then capture pattern outputs byreturnstructure order. In the above example,a/b/cwill be the 0th/1st/2nd captured tensor, andMatMul(a, b) + coutput will be the 3rd captured tensor inmatch_result.- Do not directly iterate over
inputs, because input count is not predetermined. - One
@patternmethod represents one pattern. - Multiple topologies need multiple
@patternmethods. @patterncannot be used together withpatterns(self).
3.4 Use meet_requirements for Condition Filtering (Optional)
If topology matching needs additional checks for dtype, shape, attributes, or constant values, implementmeet_requirements:
def meet_requirements(self, match_result): for node in match_result.get_matched_nodes(): if node.type == "Const": return _is_zero(node.get_attr("value")) return Falsematch_resultis the result of this match. It can get matched real nodes and captured tensors from the pattern. When using@pattern, visited external inputs are automatically captured by input index, andreturnpattern outputs are captured by return order; intermediate tensors not used asreturnoutputs are not automatically captured.
If only topology matching is sufficient, this method can be omitted; it returnsTrueby default.
3.5 Use replacement to Define Replacement Structure
The simplest replacement can directly return an input:
def replacement(self, inputs): return inputs[0]Can also use expressions to create new structures:
def replacement(self, inputs): a, b, c = inputs[:3] return GEMM(a, b, c, 1.0, 1.0)If replacement needs to read matched node attributes, add amatch_resultparameter:
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)4. When Not to Use @pattern
@patternfits most common topologies, but has a clear boundary: it automatically captures visited external inputs andreturnpattern outputs, but does not automatically capture intermediate tensors not returned as outputs.
Ifmeet_requirementsorreplacementneeds to read intermediate tensors not returned asreturnoutputs, e.g.,MatMuloutput, do not use@pattern. Instead, explicitly create pattern graph and callPattern.capture_tensorto mark intermediate tensors to read. If only need to read the final output returned byreturn, e.g.,Addoutput, continue using@pattern.
This approach is closer to C++:
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]))If only expressing patterns likeAdd(x, 0),MatMul + Add, prefer@pattern, code is shorter and closer to optimization logic.
5. Capture Tensor
Capture tensor allows retrieving the corresponding real tensor frommatch_resultby capture order after pattern matching.
Common uses:
- Check dtype or shape of an output tensor.
- Read original node attributes to pass to new nodes in replacement.
- Print matched location to confirm pass hits expected nodes.
Refer to capture tensor Python example.
6. More Strict Matching: PatternMatcherConfig
Default matcher mainly checks topology and operator types. If wanting to check Const values during matching phase, pass configuration in constructor:
from ge.passes import PatternMatcherConfigBuilder class PythonAddZeroConstValueMatchPass(PatternFusionPass): def __init__(self): super().__init__( PatternMatcherConfigBuilder() .enable_const_value_match() .build() ) @pattern def add_zero(self, inputs): return inputs[0] + 0.0 def replacement(self, inputs): return inputs[0]This is shorter, but Const value matching is strict, without floating-point tolerance or cross-dtype normalization. If judgment needs tolerance or more complex logic, put it inmeet_requirementsfor reliability.
Refer to PatternMatcherConfig Python example.
7. Writing DecomposePass
If the goal is "when seeing a certain single operator, decompose it into a set of operators", useDecomposePass.
Skeleton is as follows:
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 replacement graph composed of basic operators ...op_typesdetermines which types of nodes GE will pass to this pass.meet_requirementsthen determines which of these nodes really need replacement.
Complete example see DecomposePass Python example.
8. Running Python pass
8.1 Setting Environment
First set CANN environment variables:
source ${ASCEND_PATH}/set_env.shASCEND_PATHpoints to CANN Toolkit installation directory, more installation path information see Quick Install. Python pass runtime will load precompiled binary components built based onpybind11, which is related to Python version. CANN package contains precompiled artifacts for multiple Python versions, and defaults to installing artifacts corresponding to current Python version. Runtime will prioritize loading artifacts matching current Python version; if no matching artifacts exist, will enter fallback compilation process, fallback compilation depends onpybind11already installed in current Python environment.
Then tell GE where to load Python pass from:
export ASCEND_GE_PY_PASS_PATH=/path/to/my_pass.pyCan also point to directory:
export ASCEND_GE_PY_PASS_PATH=/path/to/pass_dir/Multiple paths separated by colon:
export ASCEND_GE_PY_PASS_PATH=/path/to/a.py:/path/to/pass_dir/Detailed scanning rules see ASCEND_GE_PY_PASS_PATH.
8.2 Offline Compilation
Offline scenario suggests usingpyatcto trigger compilation.pyatcandatccommand line parameters are consistent, but will run in current Python interpreter process, convenient for loading Python pass.
pyatc --model=./model.onnx --framework=5 --soc_version=xxx --output=./model8.3 Online Scenario
In online scenario, setASCEND_GE_PY_PASS_PATHbefore triggering GE compilation. Examples usually trigger online compilation and execution throughtorch_forward.py.
9. Verification and Troubleshooting
Recommend enabling graph dump for every development:
export DUMP_GE_GRAPH=1Then compare.pbtxtbefore and after replacement:
PreRunBegin: Before pass execution.RunCustomPass...: After custom pass execution.
If not matched, troubleshoot in this order:
| Phenomenon | Possible Cause | Check Method |
|---|---|---|
| Python file not loaded | ASCEND_GE_PY_PASS_PATHnot set, path does not exist, suffix is not.py | First confirm environment variable and path |
| Class loaded but pass not executed | No registration decorator used, or registration stage incorrect | Check@register_fusion_pass/@register_decompose_pass |
| Pattern not matched | Operator type, input count or output boundary inconsistent | Compare real topology in dump graph |
| Matched but not replaced | meet_requirementsreturnedFalse | Print matched node attributes |
| Graph abnormal after replacement | replacement output did not cover Tensor needed by external consumers | Go back to mechanism document to check boundary rules |
When more logs needed, can set:
export ASCEND_SLOG_PRINT_TO_STDOUT=1 export ASCEND_GLOBAL_LOG_LEVEL=0When usingpyatc, can also add--log=debug.
10. Recommended Reading Order
- Fusion Pattern Pass Mechanism
- AddZeroPass Python example
- MatMul+Add Python example
- capture tensor Python example
- PatternMatcherConfig Python example
- DecomposePass Python example
【免费下载链接】geGE(Graph Engine)是面向昇腾的图编译器和执行器,提供了计算图优化、多流并行、内存复用和模型下沉等技术手段,加速模型执行效率,减少模型内存占用。 GE 提供对 PyTorch、TensorFlow 前端的友好接入能力,并同时支持 onnx、pb 等主流模型格式的解析与编译。项目地址: https://gitcode.com/cann/ge
创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考