Coverage for /opt/cloud/slavespace/usr1/096471637100f3de0fcfc01072822a80/dttest/api/python/ge/ge/passes/base.py: 94%

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# ----------------------------------------------------------------------------------------------------------- 

12 

13"""Base definitions for Python GE passes.""" 

14 

15from __future__ import annotations 

16 

17from enum import Enum 

18from typing import TYPE_CHECKING, Iterable, List, Optional, Union 

19 

20from ._native import ( # noqa: F401 

21 MatchResult, 

22 PassContext, 

23 PatternMatcherConfig, 

24 PatternMatcherConfigBuilder, 

25 SubgraphBoundary, 

26 SubgraphInput, 

27 SubgraphOutput, 

28 SubgraphRewriter, 

29) 

30from ._native import infer_shape as _infer_shape 

31 

32if TYPE_CHECKING: 

33 from ge.graph.graph import Graph 

34 from ge.graph.node import Node 

35 

36 from .pattern import Pattern 

37 

38 

39PatternOrGraph = Union["Pattern", "Graph"] 

40StatusLike = Optional[Union[bool, int]] 

41 

42 

43def infer_shape( 

44 replacement: "Graph", source: Union[MatchResult, "Node", SubgraphBoundary] 

45) -> None: 

46 """Infer shape, data type, and format for a replacement graph from a matched source.""" 

47 from ge.graph import Graph, Node 

48 

49 if not isinstance(replacement, Graph): 

50 raise TypeError("replacement must be a ge.graph.Graph") 

51 if not isinstance(source, (MatchResult, Node, SubgraphBoundary)): 

52 raise TypeError("source must be MatchResult, Node, or SubgraphBoundary") 

53 if not replacement.get_all_nodes(): 

54 raise RuntimeError("replacement graph is empty") 

55 _infer_shape(replacement, source) 

56 

57 

58class PassStage(str, Enum): 

59 """Python-facing pass stage names.""" 

60 

61 BEFORE_INFER_SHAPE = "BeforeInferShape" 

62 AFTER_INFER_SHAPE = "AfterInferShape" 

63 AFTER_BUILTIN_FUSION_PASS = "AfterBuiltinFusionPass" 

64 AFTER_ORIGIN_GRAPH_OPTIMIZE = "AfterOriginGraphOptimize" 

65 

66 

67class FusionBasePass: 

68 """Python FusionBasePass contract.""" 

69 

70 def run(self, graph: Graph, context: PassContext) -> StatusLike: 

71 raise NotImplementedError("FusionBasePass.run must be implemented") 

72 

73 

74class PatternFusionPass(FusionBasePass): 

75 """Python PatternFusionPass contract. 

76 

77 The execution engine calls ``patterns()``, ``meet_requirements()``, and 

78 ``replacement()`` — **not** ``run()``. Overriding ``run()`` in a 

79 ``PatternFusionPass`` subclass has no effect; implement the three hook 

80 methods above instead. 

81 

82 Subclasses can use either the legacy graph-building hooks: 

83 - ``patterns(self) -> Iterable[Pattern | Graph]`` 

84 - ``meet_requirements(self, match_result) -> bool`` 

85 - ``replacement(self, match_result) -> Graph`` 

86 - the same hooks may append ``context: PassContext`` 

87 

88 or expression-style pattern methods: 

89 - ``@pattern def name(self, inputs) -> TensorHolder | list[TensorHolder] | tuple[TensorHolder, ...]`` 

90 - ``replacement(self, inputs) -> TensorHolder | list[TensorHolder] | tuple[TensorHolder, ...] | Graph`` 

91 - ``replacement(self, inputs, match_result)`` when match details are needed 

92 - ``replacement(self, inputs, context)`` or ``replacement(self, inputs, match_result, context)`` 

93 

94 In a ``@pattern`` method, a list/tuple means multiple outputs of one 

95 pattern, not multiple patterns. Declare several ``@pattern`` methods for 

96 a multiple-pattern pass. Expression-style patterns automatically capture 

97 accessed inputs in input-index order, then returned outputs in return order. 

98 """ 

99 

100 def __init__(self, matcher_config: Optional[PatternMatcherConfig] = None) -> None: 

101 self._matcher_config = matcher_config 

102 

103 def __init_subclass__(cls, **kwargs) -> None: 

104 super().__init_subclass__(**kwargs) 

105 if "run" in cls.__dict__: 

106 raise TypeError( 

107 f"{cls.__name__} overrides run(), which is never invoked " 

108 f"by the PatternFusionPass execution path. " 

109 f"Implement patterns()/replacement() instead." 

110 ) 

111 from .pattern import ( 

112 _adapt_context_hook, 

113 _adapt_decorated_pattern_methods, 

114 _adapt_replacement_hook, 

115 _has_decorated_pattern_methods, 

116 ) 

117 

118 if _has_decorated_pattern_methods(cls): 

119 if "patterns" in cls.__dict__: 

120 raise TypeError( 

121 f"{cls.__name__} cannot combine @pattern methods with patterns(). " 

122 f"Use one style for declaring patterns." 

123 ) 

124 cls.patterns = _adapt_decorated_pattern_methods(cls) 

125 if "meet_requirements" in cls.__dict__: 

126 cls.meet_requirements = _adapt_context_hook( 

127 cls.__dict__["meet_requirements"], 

128 "meet_requirements", 

129 "match_result", 

130 ) 

131 if "replacement" in cls.__dict__: 

132 cls.replacement = _adapt_replacement_hook(cls.__dict__["replacement"]) 

133 

134 @property 

135 def matcher_config(self) -> Optional[PatternMatcherConfig]: 

136 return self._matcher_config 

137 

138 def patterns(self) -> Iterable[PatternOrGraph]: 

139 raise NotImplementedError("PatternFusionPass.patterns must be implemented") 

140 

141 def meet_requirements( 

142 self, match_result: MatchResult, context: Optional[PassContext] = None 

143 ) -> bool: 

144 return True 

145 

146 def replacement( 

147 self, match_result: MatchResult, context: Optional[PassContext] = None 

148 ) -> "Graph": 

149 raise NotImplementedError("PatternFusionPass.replacement must be implemented") 

150 

151 

152class DecomposePass(FusionBasePass): 

153 """Python DecomposePass contract. 

154 

155 The execution engine calls ``meet_requirements()`` and ``replacement()`` 

156 for matched nodes — **not** ``run()``. Overriding ``run()`` in a 

157 ``DecomposePass`` subclass has no effect; implement the two hook methods 

158 above instead. 

159 """ 

160 

161 op_types: Optional[List[str]] = None 

162 

163 def __init_subclass__(cls, **kwargs) -> None: 

164 super().__init_subclass__(**kwargs) 

165 if "run" in cls.__dict__: 

166 raise TypeError( 

167 f"{cls.__name__} overrides run(), which is never invoked " 

168 f"by the DecomposePass execution path. " 

169 f"Implement meet_requirements()/replacement() instead." 

170 ) 

171 from .pattern import _adapt_context_hook 

172 

173 if "meet_requirements" in cls.__dict__: 

174 cls.meet_requirements = _adapt_context_hook( 

175 cls.__dict__["meet_requirements"], "meet_requirements", "node" 

176 ) 

177 if "replacement" in cls.__dict__: 

178 cls.replacement = _adapt_context_hook( 

179 cls.__dict__["replacement"], "replacement", "node" 

180 ) 

181 

182 def meet_requirements( 

183 self, node: "Node", context: Optional[PassContext] = None 

184 ) -> bool: 

185 return True 

186 

187 def replacement( 

188 self, node: "Node", context: Optional[PassContext] = None 

189 ) -> "Graph": 

190 raise NotImplementedError("DecomposePass.replacement must be implemented")