Coverage for /opt/cloud/slavespace/usr1/096471637100f3de0fcfc01072822a80/dttest/api/python/ge/ge/passes/pattern.py: 91%
163 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-04 11:36 +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__(
70 self, index: Union[int, slice]
71 ) -> Union[TensorHolder, List[TensorHolder]]:
72 if isinstance(index, slice):
73 return self._get_slice(index)
74 if not isinstance(index, int):
75 raise TypeError("PatternInputs index must be an integer or slice")
76 if index < 0:
77 raise ValueError("PatternInputs index must be non-negative")
78 self._ensure_created(index)
79 return self._inputs[index]
81 def __iter__(self) -> Iterator[TensorHolder]:
82 raise TypeError(
83 "PatternInputs cannot be iterated because the input count is unknown; use inputs[:N]"
84 )
86 def _get_slice(self, index_slice: slice) -> List[TensorHolder]:
87 if index_slice.step not in (None, 1):
88 raise ValueError("PatternInputs slice step must be 1")
89 start = 0 if index_slice.start is None else index_slice.start
90 stop = index_slice.stop
91 if stop is None:
92 raise ValueError(
93 "PatternInputs slice must provide a stop index, for example inputs[:3]"
94 )
95 if start < 0 or stop < 0:
96 raise ValueError("PatternInputs slice indices must be non-negative")
97 if stop < start:
98 raise ValueError(
99 "PatternInputs slice stop must be greater than or equal to start"
100 )
101 return [self[i] for i in range(start, stop)]
103 def _ensure_created(self, index: int) -> None:
104 for input_index in range(len(self._inputs), index + 1):
105 self._inputs[input_index] = self._builder.create_input(input_index)
107 def created(self) -> List[TensorHolder]:
108 """Return created inputs in graph input index order."""
110 return [self._inputs[index] for index in sorted(self._inputs)]
113def create_pattern(graph: Graph) -> "Pattern":
114 """Build a native Pattern from a pattern graph."""
115 from ._native import Pattern
117 return Pattern(graph)
120def ensure_pattern(pattern_or_graph: Union["Pattern", Graph]) -> "Pattern":
121 """Convert a Graph into Pattern on demand for bridge helpers."""
122 from ._native import Pattern
124 if isinstance(pattern_or_graph, Graph):
125 return create_pattern(pattern_or_graph)
126 if isinstance(pattern_or_graph, Pattern):
127 return pattern_or_graph
128 raise TypeError("PatternFusionPass.patterns must return Pattern or Graph objects")
131def pattern(method: Callable[..., object]) -> Callable[..., object]:
132 """Mark a ``PatternFusionPass`` method as one expression-style pattern."""
134 setattr(method, _PATTERN_METHOD_MARK, True)
135 return method
138def _has_decorated_pattern_methods(cls: Type[object]) -> bool:
139 return bool(_get_decorated_pattern_methods(cls))
142def _adapt_decorated_pattern_methods(cls: Type[object]) -> Callable[..., object]:
143 """Build legacy ``patterns(self)`` from ``@pattern`` expression methods."""
145 methods = _get_decorated_pattern_methods(cls)
146 for method in methods:
147 _check_required_arg_count(
148 method,
149 min_count=2,
150 max_count=2,
151 message="@pattern methods only support method(self, inputs)",
152 )
154 def wrapper(self) -> List["Pattern"]:
155 patterns = []
156 for method in methods:
157 builder = GraphBuilder(_decorated_pattern_graph_name(self, method))
158 inputs = PatternInputs(builder)
159 result = method(self, inputs)
160 patterns.extend(
161 _build_patterns_from_expression_result(builder, inputs, result)
162 )
163 return patterns
165 return wrapper
168def _adapt_expression_replacement(
169 method: Callable[..., object],
170) -> Callable[..., object]:
171 """Wrap expression replacement hooks into ``replacement(match_result, context)``."""
173 positional_params = _positional_params(method)
174 if len(positional_params) < 2 or positional_params[1].name != "inputs":
175 return method
176 _check_required_arg_count(
177 method,
178 min_count=2,
179 max_count=4,
180 message="Expression-style replacement only supports replacement(self, inputs), "
181 "replacement(self, inputs, match_result), replacement(self, inputs, context), or "
182 "replacement(self, inputs, match_result, context)",
183 )
184 positional_count = _positional_arg_count(method)
185 accepts_match_result = positional_count in (3, 4) and not (
186 positional_count == 3 and positional_params[2].name == "context"
187 )
188 accepts_context = positional_count in (3, 4) and (
189 positional_count == 4 or positional_params[2].name == "context"
190 )
191 if positional_count == 4 and positional_params[3].name != "context":
192 raise TypeError(
193 "Expression-style replacement with four arguments must use "
194 "replacement(self, inputs, match_result, context)"
195 )
197 @wraps(method)
198 def wrapper(self, match_result: object, context: object = None) -> Graph:
199 builder = GraphBuilder(_default_graph_name(self, "replacement"))
200 inputs = PatternInputs(builder)
201 if accepts_match_result and accepts_context:
202 result = method(self, inputs, match_result, context)
203 elif accepts_match_result:
204 result = method(self, inputs, match_result)
205 elif accepts_context:
206 result = method(self, inputs, context)
207 else:
208 result = method(self, inputs)
209 return _build_replacement_from_expression_result(builder, result)
211 return wrapper
214def _adapt_context_hook(
215 method: Callable[..., object], hook_name: str, value_name: str
216) -> Callable[..., object]:
217 """Normalize a graph-building hook to ``(value, context)``."""
219 _check_required_arg_count(
220 method,
221 min_count=2,
222 max_count=3,
223 message=f"{hook_name} only supports {hook_name}(self, {value_name}) or "
224 f"{hook_name}(self, {value_name}, context)",
225 )
226 positional_count = _positional_arg_count(method)
227 accepts_context = positional_count == 3
229 @wraps(method)
230 def wrapper(self, value: object, context: object = None):
231 if accepts_context:
232 return method(self, value, context)
233 return method(self, value)
235 return wrapper
238def _adapt_replacement_hook(
239 method: Callable[..., object],
240) -> Callable[..., object]:
241 """Normalize an expression-style or graph-building replacement hook."""
243 adapted = _adapt_expression_replacement(method)
244 if adapted is not method:
245 return adapted
246 return _adapt_context_hook(method, "replacement", "match_result")
249def _check_required_arg_count(
250 method: Callable[..., object], *, min_count: int, max_count: int, message: str
251) -> None:
252 count = _positional_arg_count(method)
253 if (
254 count < min_count
255 or count > max_count
256 or _has_required_keyword_only_args(method)
257 ):
258 raise TypeError(message)
261def _positional_arg_count(method: Callable[..., object]) -> int:
262 return len(_positional_params(method))
265def _positional_params(method: Callable[..., object]) -> List[inspect.Parameter]:
266 signature = inspect.signature(method)
267 return [
268 param
269 for param in signature.parameters.values()
270 if param.kind
271 in (inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD)
272 ]
275def _has_required_keyword_only_args(method: Callable[..., object]) -> bool:
276 signature = inspect.signature(method)
277 return any(
278 param.kind == inspect.Parameter.KEYWORD_ONLY
279 and param.default is inspect.Parameter.empty
280 for param in signature.parameters.values()
281 )
284def _get_decorated_pattern_methods(cls: Type[object]) -> List[Callable[..., object]]:
285 return [
286 method
287 for method in cls.__dict__.values()
288 if callable(method) and getattr(method, _PATTERN_METHOD_MARK, False)
289 ]
292def _default_graph_name(instance: object, suffix: str) -> str:
293 return f"{instance.__class__.__name__}_{suffix}"
296def _decorated_pattern_graph_name(
297 instance: object, method: Callable[..., object]
298) -> str:
299 return f"{instance.__class__.__name__}_{method.__name__}_pattern"
302def _build_patterns_from_expression_result(
303 builder: GraphBuilder, inputs: PatternInputs, result: object
304) -> List["Pattern"]:
305 if _is_pattern_or_graph(result):
306 return [ensure_pattern(result)]
307 if (
308 isinstance(result, (list, tuple))
309 and result
310 and all(_is_pattern_or_graph(item) for item in result)
311 ):
312 raise TypeError(
313 "A @pattern method supports a single pattern only. "
314 "For multiple patterns, declare several @pattern methods or use legacy patterns(self)."
315 )
317 pattern_outputs = _normalize_tensor_outputs(result, "patterns")
318 graph = builder.build_and_reset(pattern_outputs)
319 built_pattern = create_pattern(graph)
320 for input_tensor in inputs.created():
321 built_pattern.capture_tensor(input_tensor)
322 for output_tensor in pattern_outputs:
323 built_pattern.capture_tensor(output_tensor)
324 return [built_pattern]
327def _build_replacement_from_expression_result(
328 builder: GraphBuilder, result: object
329) -> Graph:
330 if isinstance(result, Graph):
331 return result
332 return builder.build_and_reset(_normalize_tensor_outputs(result, "replacement"))
335def _normalize_tensor_outputs(result: object, hook_name: str) -> List[TensorHolder]:
336 if isinstance(result, TensorHolder):
337 return [result]
338 if isinstance(result, tuple):
339 result = list(result)
340 if (
341 isinstance(result, list)
342 and result
343 and all(isinstance(item, TensorHolder) for item in result)
344 ):
345 return result
346 raise TypeError(
347 f"Expression-style {hook_name} must return a TensorHolder or a non-empty list/tuple of TensorHolder"
348 )
351def _is_pattern_or_graph(result: object) -> bool:
352 from ._native import Pattern
354 return isinstance(result, (Pattern, Graph))