Coverage for /opt/cloud/slavespace/usr1/096471637100f3de0fcfc01072822a80/dttest/api/python/ge/ge/passes/pattern.py: 89%
140 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-27 10:02 +0800
1#!/usr/bin/env python3
2# -*- coding: utf-8 -*-
3# -----------------------------------------------------------------------------------------------------------
4# Copyright (c) 2026 Huawei Technologies Co., Ltd.
5# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
6# CANN Open Software License Agreement Version 2.0 (the "License").
7# Please refer to the License for details. You may not use this file except in compliance with the License.
8# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
9# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
10# See LICENSE in the root of the software repository for the full text of the License.
11# -----------------------------------------------------------------------------------------------------------
13"""PatternFusionPass helper types and constructors."""
15from __future__ import annotations
17import inspect
18from dataclasses import dataclass
19from functools import wraps
20from typing import TYPE_CHECKING, Callable, Dict, Iterator, List, Type, Union, overload
22from ge.es.graph_builder import GraphBuilder
23from ge.es.tensor_holder import TensorHolder
24from ge.graph import Graph, Node
26if TYPE_CHECKING:
27 from ._native import Pattern
30_PATTERN_METHOD_MARK = "__ge_expression_pattern_method__"
33def __getattr__(name: str):
34 if name == "Pattern":
35 from ._native import Pattern
37 globals()["Pattern"] = Pattern
38 return Pattern
39 raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
42@dataclass(frozen=True)
43class NodeIo:
44 """Python helper describing one node output."""
46 node: Node
47 index: int = 0
50class PatternInputs:
51 """Lazy graph input collection passed to expression-style pass hooks.
53 Users receive this object from ``@pattern`` methods or
54 ``replacement(self, inputs)``. Accessing ``inputs[i]`` creates graph input
55 ``i`` on demand. Use ``inputs[:N]`` to declare multiple inputs explicitly.
56 The input count is intentionally unknown, so direct iteration is rejected.
57 """
59 def __init__(self, builder: GraphBuilder) -> None:
60 self._builder = builder
61 self._inputs: Dict[int, TensorHolder] = {}
63 @overload
64 def __getitem__(self, index: int) -> TensorHolder: ...
66 @overload
67 def __getitem__(self, index: slice) -> List[TensorHolder]: ...
69 def __getitem__(self, index: Union[int, slice]) -> Union[TensorHolder, List[TensorHolder]]:
70 if isinstance(index, slice):
71 return self._get_slice(index)
72 if not isinstance(index, int):
73 raise TypeError("PatternInputs index must be an integer or slice")
74 if index < 0:
75 raise ValueError("PatternInputs index must be non-negative")
76 self._ensure_created(index)
77 return self._inputs[index]
79 def __iter__(self) -> Iterator[TensorHolder]:
80 raise TypeError("PatternInputs cannot be iterated because the input count is unknown; use inputs[:N]")
82 def _get_slice(self, index_slice: slice) -> List[TensorHolder]:
83 if index_slice.step not in (None, 1):
84 raise ValueError("PatternInputs slice step must be 1")
85 start = 0 if index_slice.start is None else index_slice.start
86 stop = index_slice.stop
87 if stop is None:
88 raise ValueError("PatternInputs slice must provide a stop index, for example inputs[:3]")
89 if start < 0 or stop < 0:
90 raise ValueError("PatternInputs slice indices must be non-negative")
91 if stop < start:
92 raise ValueError("PatternInputs slice stop must be greater than or equal to start")
93 return [self[i] for i in range(start, stop)]
95 def _ensure_created(self, index: int) -> None:
96 for input_index in range(len(self._inputs), index + 1):
97 self._inputs[input_index] = self._builder.create_input(input_index)
99 def created(self) -> List[TensorHolder]:
100 """Return created inputs in graph input index order."""
102 return [self._inputs[index] for index in sorted(self._inputs)]
105def create_pattern(graph: Graph) -> "Pattern":
106 """Build a native Pattern from a pattern graph."""
107 from ._native import Pattern
109 return Pattern(graph)
112def ensure_pattern(pattern_or_graph: Union["Pattern", Graph]) -> "Pattern":
113 """Convert a Graph into Pattern on demand for bridge helpers."""
114 from ._native import Pattern
116 if isinstance(pattern_or_graph, Graph):
117 return create_pattern(pattern_or_graph)
118 if isinstance(pattern_or_graph, Pattern):
119 return pattern_or_graph
120 raise TypeError("PatternFusionPass.patterns must return Pattern or Graph objects")
123def pattern(method: Callable[..., object]) -> Callable[..., object]:
124 """Mark a ``PatternFusionPass`` method as one expression-style pattern."""
126 setattr(method, _PATTERN_METHOD_MARK, True)
127 return method
130def _has_decorated_pattern_methods(cls: Type[object]) -> bool:
131 return bool(_get_decorated_pattern_methods(cls))
134def _adapt_decorated_pattern_methods(cls: Type[object]) -> Callable[..., object]:
135 """Build legacy ``patterns(self)`` from ``@pattern`` expression methods."""
137 methods = _get_decorated_pattern_methods(cls)
138 for method in methods:
139 _check_required_arg_count(
140 method,
141 min_count=2,
142 max_count=2,
143 message="@pattern methods only support method(self, inputs)",
144 )
146 def wrapper(self) -> List["Pattern"]:
147 patterns = []
148 for method in methods:
149 builder = GraphBuilder(_decorated_pattern_graph_name(self, method))
150 inputs = PatternInputs(builder)
151 result = method(self, inputs)
152 patterns.extend(_build_patterns_from_expression_result(builder, inputs, result))
153 return patterns
155 return wrapper
158def _adapt_expression_replacement(
159 method: Callable[..., object],
160) -> Callable[..., object]:
161 """Wrap ``replacement(self, inputs[, match_result])`` into the legacy hook."""
163 positional_params = _positional_params(method)
164 if len(positional_params) < 2 or positional_params[1].name != "inputs":
165 return method
166 _check_required_arg_count(
167 method,
168 min_count=2,
169 max_count=3,
170 message="Expression-style replacement only supports "
171 "replacement(self, inputs) or replacement(self, inputs, match_result)",
172 )
173 accepts_match_result = _positional_arg_count(method) == 3
175 @wraps(method)
176 def wrapper(self, match_result: object) -> Graph:
177 builder = GraphBuilder(_default_graph_name(self, "replacement"))
178 inputs = PatternInputs(builder)
179 if accepts_match_result:
180 result = method(self, inputs, match_result)
181 else:
182 result = method(self, inputs)
183 return _build_replacement_from_expression_result(builder, result)
185 return wrapper
188def _check_required_arg_count(method: Callable[..., object], *, min_count: int, max_count: int, message: str) -> None:
189 count = _positional_arg_count(method)
190 if count < min_count or count > max_count or _has_required_keyword_only_args(method):
191 raise TypeError(message)
194def _positional_arg_count(method: Callable[..., object]) -> int:
195 return len(_positional_params(method))
198def _positional_params(method: Callable[..., object]) -> List[inspect.Parameter]:
199 signature = inspect.signature(method)
200 return [
201 param
202 for param in signature.parameters.values()
203 if param.kind in (inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD)
204 ]
207def _has_required_keyword_only_args(method: Callable[..., object]) -> bool:
208 signature = inspect.signature(method)
209 return any(
210 param.kind == inspect.Parameter.KEYWORD_ONLY and param.default is inspect.Parameter.empty
211 for param in signature.parameters.values()
212 )
215def _get_decorated_pattern_methods(cls: Type[object]) -> List[Callable[..., object]]:
216 return [
217 method for method in cls.__dict__.values() if callable(method) and getattr(method, _PATTERN_METHOD_MARK, False)
218 ]
221def _default_graph_name(instance: object, suffix: str) -> str:
222 return f"{instance.__class__.__name__}_{suffix}"
225def _decorated_pattern_graph_name(instance: object, method: Callable[..., object]) -> str:
226 return f"{instance.__class__.__name__}_{method.__name__}_pattern"
229def _build_patterns_from_expression_result(
230 builder: GraphBuilder, inputs: PatternInputs, result: object
231) -> List["Pattern"]:
232 if _is_pattern_or_graph(result):
233 return [ensure_pattern(result)]
234 if isinstance(result, (list, tuple)) and result and all(_is_pattern_or_graph(item) for item in result):
235 raise TypeError(
236 "A @pattern method supports a single pattern only. "
237 "For multiple patterns, declare several @pattern methods or use legacy patterns(self)."
238 )
240 pattern_outputs = _normalize_tensor_outputs(result, "patterns")
241 graph = builder.build_and_reset(pattern_outputs)
242 built_pattern = create_pattern(graph)
243 for input_tensor in inputs.created():
244 built_pattern.capture_tensor(input_tensor)
245 for output_tensor in pattern_outputs:
246 built_pattern.capture_tensor(output_tensor)
247 return [built_pattern]
250def _build_replacement_from_expression_result(builder: GraphBuilder, result: object) -> Graph:
251 if isinstance(result, Graph):
252 return result
253 return builder.build_and_reset(_normalize_tensor_outputs(result, "replacement"))
256def _normalize_tensor_outputs(result: object, hook_name: str) -> List[TensorHolder]:
257 if isinstance(result, TensorHolder):
258 return [result]
259 if isinstance(result, tuple):
260 result = list(result)
261 if isinstance(result, list) and result and all(isinstance(item, TensorHolder) for item in result):
262 return result
263 raise TypeError(
264 f"Expression-style {hook_name} must return a TensorHolder or a non-empty list/tuple of TensorHolder"
265 )
268def _is_pattern_or_graph(result: object) -> bool:
269 from ._native import Pattern
271 return isinstance(result, (Pattern, Graph))