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

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"""Bridge-facing Python runtime helpers for GE passes.""" 

14 

15import sys 

16import threading 

17from collections.abc import Iterable 

18from dataclasses import dataclass 

19from typing import Dict, List, Optional, cast 

20 

21from ge.graph import Graph, Node 

22 

23from ._native import ( 

24 borrow_match_result, 

25 borrow_node, 

26 clone_pattern_matcher_config, 

27 release_graph, 

28) 

29from .base import ( 

30 DecomposePass, 

31 FusionBasePass, 

32 PassContext, 

33 PatternFusionPass, 

34 PatternMatcherConfig, 

35 StatusLike, 

36) 

37from .bootstrap import get_registered_passes, load_pass_plugins 

38from .pattern import ensure_pattern 

39from .registry import get_registered_pass_by_descriptor_key 

40 

41 

42@dataclass 

43class _PassHolder: 

44 descriptor_key: str 

45 instance_id: str 

46 instance: FusionBasePass 

47 

48 

49_HOLDER_LOCK = threading.RLock() 

50_PASS_HOLDERS: Dict[str, _PassHolder] = {} 

51 

52 

53def load_and_get_pass_descriptors() -> list: 

54 load_pass_plugins() 

55 return get_registered_passes() 

56 

57 

58def _get_holder(instance_id: str) -> _PassHolder: 

59 with _HOLDER_LOCK: 

60 holder = _PASS_HOLDERS.get(instance_id) 

61 if holder is None: 

62 raise KeyError(f"python pass holder is not created: {instance_id}") 

63 return holder 

64 

65 

66def _get_fusion_base_pass(instance_id: str) -> FusionBasePass: 

67 return _get_holder(instance_id).instance 

68 

69 

70def _get_pattern_fusion_pass(instance_id: str) -> PatternFusionPass: 

71 instance = _get_holder(instance_id).instance 

72 if not isinstance(instance, PatternFusionPass): 

73 raise TypeError(f"python pass holder is not PatternFusionPass: {instance_id}") 

74 return instance 

75 

76 

77def _get_decompose_pass(instance_id: str) -> DecomposePass: 

78 instance = _get_holder(instance_id).instance 

79 if not isinstance(instance, DecomposePass): 

80 raise TypeError(f"python pass holder is not DecomposePass: {instance_id}") 

81 return instance 

82 

83 

84def create_pass_holder(instance_id: str, descriptor_key: str) -> bool: 

85 descriptor = get_registered_pass_by_descriptor_key(descriptor_key) 

86 if descriptor is None: 

87 raise KeyError(f"python pass descriptor_key not found: {descriptor_key}") 

88 with _HOLDER_LOCK: 

89 if instance_id in _PASS_HOLDERS: 

90 return True 

91 _PASS_HOLDERS[instance_id] = _PassHolder( 

92 descriptor_key=descriptor_key, 

93 instance_id=instance_id, 

94 instance=descriptor.cls(), 

95 ) 

96 return True 

97 

98 

99def destroy_pass_holder(instance_id: str) -> bool: 

100 with _HOLDER_LOCK: 

101 return _PASS_HOLDERS.pop(instance_id, None) is not None 

102 

103 

104def run_fusion_base_pass( 

105 instance_id: str, graph: Graph, context: Optional[PassContext] = None 

106) -> StatusLike: 

107 pass_instance = _get_fusion_base_pass(instance_id) 

108 if context is not None and not isinstance(context, PassContext): 

109 raise TypeError("context type error") 

110 return pass_instance.run(graph, cast(PassContext, context)) 

111 

112 

113def _release_replacement_graph(replacement: Graph, pass_name: str) -> int: 

114 if not isinstance(replacement, Graph): 

115 raise TypeError(f"{pass_name}.replacement must return ge.graph.Graph") 

116 return release_graph(replacement) 

117 

118 

119def get_pass_patterns(instance_id: str) -> List[int]: 

120 pass_instance = _get_pattern_fusion_pass(instance_id) 

121 patterns = pass_instance.patterns() 

122 if patterns is None: 

123 return [] 

124 if not isinstance(patterns, Iterable) or isinstance(patterns, (str, bytes)): 

125 raise TypeError( 

126 "PatternFusionPass.patterns must return an iterable of Pattern or Graph" 

127 ) 

128 

129 released_patterns = [] 

130 for item in patterns: 

131 pattern = ensure_pattern(item) 

132 released_patterns.append(pattern.release()) 

133 return released_patterns 

134 

135 

136def get_pattern_matcher_config(instance_id: str) -> Optional[int]: 

137 pass_instance = _get_pattern_fusion_pass(instance_id) 

138 matcher_config = pass_instance.matcher_config 

139 if matcher_config is None: 

140 return None 

141 if not isinstance(matcher_config, PatternMatcherConfig): 

142 raise TypeError( 

143 "PatternFusionPass.matcher_config must be PatternMatcherConfig or None" 

144 ) 

145 return clone_pattern_matcher_config(matcher_config) 

146 

147 

148def call_meet_requirements( 

149 instance_id: str, match_result_handle: int, context: Optional[PassContext] = None 

150) -> bool: 

151 pass_instance = _get_pattern_fusion_pass(instance_id) 

152 match_result = borrow_match_result(match_result_handle) 

153 try: 

154 return bool(pass_instance.meet_requirements(match_result, context)) 

155 finally: 

156 match_result._invalidate() 

157 

158 

159def call_replacement( 

160 instance_id: str, match_result_handle: int, context: Optional[PassContext] = None 

161) -> int: 

162 pass_instance = _get_pattern_fusion_pass(instance_id) 

163 match_result = borrow_match_result(match_result_handle) 

164 try: 

165 replacement = pass_instance.replacement(match_result, context) 

166 finally: 

167 match_result._invalidate() 

168 

169 return _release_replacement_graph(replacement, "PatternFusionPass") 

170 

171 

172def call_decompose_meet_requirements( 

173 instance_id: str, node_handle: int, context: Optional[PassContext] = None 

174) -> bool: 

175 pass_instance = _get_decompose_pass(instance_id) 

176 node = cast(Node, borrow_node(node_handle)) 

177 return bool(pass_instance.meet_requirements(node, context)) 

178 

179 

180def call_decompose_replacement( 

181 instance_id: str, node_handle: int, context: Optional[PassContext] = None 

182) -> int: 

183 pass_instance = _get_decompose_pass(instance_id) 

184 node = cast(Node, borrow_node(node_handle)) 

185 replacement = pass_instance.replacement(node, context) 

186 return _release_replacement_graph(replacement, "DecomposePass") 

187 

188 

189def clear_pass_holders() -> None: 

190 with _HOLDER_LOCK: 

191 _PASS_HOLDERS.clear() 

192 

193 

194def clear_loaded_pass_modules() -> None: 

195 """Clear all dynamically loaded pass modules from sys.modules to avoid test pollution.""" 

196 keys_to_remove = [key for key in sys.modules if key.startswith("_ge_py_pass_")] 

197 for key in keys_to_remove: 

198 del sys.modules[key]