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 borrow_match_result, borrow_node, clone_pattern_matcher_config, release_graph 

24from .base import ( 

25 DecomposePass, 

26 FusionBasePass, 

27 PassContext, 

28 PatternFusionPass, 

29 PatternMatcherConfig, 

30 StatusLike, 

31) 

32from .bootstrap import get_registered_passes, load_pass_plugins 

33from .pattern import ensure_pattern 

34from .registry import get_registered_pass_by_descriptor_key 

35 

36 

37@dataclass 

38class _PassHolder: 

39 descriptor_key: str 

40 instance_id: str 

41 instance: FusionBasePass 

42 

43 

44_HOLDER_LOCK = threading.RLock() 

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

46 

47 

48def load_and_get_pass_descriptors() -> list: 

49 load_pass_plugins() 

50 return get_registered_passes() 

51 

52 

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

54 with _HOLDER_LOCK: 

55 holder = _PASS_HOLDERS.get(instance_id) 

56 if holder is None: 

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

58 return holder 

59 

60 

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

62 return _get_holder(instance_id).instance 

63 

64 

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

66 instance = _get_holder(instance_id).instance 

67 if not isinstance(instance, PatternFusionPass): 

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

69 return instance 

70 

71 

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

73 instance = _get_holder(instance_id).instance 

74 if not isinstance(instance, DecomposePass): 

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

76 return instance 

77 

78 

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

80 descriptor = get_registered_pass_by_descriptor_key(descriptor_key) 

81 if descriptor is None: 

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

83 with _HOLDER_LOCK: 

84 if instance_id in _PASS_HOLDERS: 

85 return True 

86 _PASS_HOLDERS[instance_id] = _PassHolder( 

87 descriptor_key=descriptor_key, 

88 instance_id=instance_id, 

89 instance=descriptor.cls(), 

90 ) 

91 return True 

92 

93 

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

95 with _HOLDER_LOCK: 

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

97 

98 

99def run_fusion_base_pass(instance_id: str, graph: Graph, context: Optional[PassContext] = None) -> StatusLike: 

100 pass_instance = _get_fusion_base_pass(instance_id) 

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

102 raise TypeError("context type error") 

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

104 

105 

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

107 if not isinstance(replacement, Graph): 

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

109 return release_graph(replacement) 

110 

111 

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

113 pass_instance = _get_pattern_fusion_pass(instance_id) 

114 patterns = pass_instance.patterns() 

115 if patterns is None: 

116 return [] 

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

118 raise TypeError("PatternFusionPass.patterns must return an iterable of Pattern or Graph") 

119 

120 released_patterns = [] 

121 for item in patterns: 

122 pattern = ensure_pattern(item) 

123 released_patterns.append(pattern.release()) 

124 return released_patterns 

125 

126 

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

128 pass_instance = _get_pattern_fusion_pass(instance_id) 

129 matcher_config = pass_instance.matcher_config 

130 if matcher_config is None: 

131 return None 

132 if not isinstance(matcher_config, PatternMatcherConfig): 

133 raise TypeError("PatternFusionPass.matcher_config must be PatternMatcherConfig or None") 

134 return clone_pattern_matcher_config(matcher_config) 

135 

136 

137def call_meet_requirements(instance_id: str, match_result_handle: int) -> bool: 

138 pass_instance = _get_pattern_fusion_pass(instance_id) 

139 match_result = borrow_match_result(match_result_handle) 

140 try: 

141 return bool(pass_instance.meet_requirements(match_result)) 

142 finally: 

143 match_result._invalidate() 

144 

145 

146def call_replacement(instance_id: str, match_result_handle: int) -> int: 

147 pass_instance = _get_pattern_fusion_pass(instance_id) 

148 match_result = borrow_match_result(match_result_handle) 

149 try: 

150 replacement = pass_instance.replacement(match_result) 

151 finally: 

152 match_result._invalidate() 

153 

154 return _release_replacement_graph(replacement, "PatternFusionPass") 

155 

156 

157def call_decompose_meet_requirements(instance_id: str, node_handle: int) -> bool: 

158 pass_instance = _get_decompose_pass(instance_id) 

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

160 return bool(pass_instance.meet_requirements(node)) 

161 

162 

163def call_decompose_replacement(instance_id: str, node_handle: int) -> int: 

164 pass_instance = _get_decompose_pass(instance_id) 

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

166 replacement = pass_instance.replacement(node) 

167 return _release_replacement_graph(replacement, "DecomposePass") 

168 

169 

170def clear_pass_holders() -> None: 

171 with _HOLDER_LOCK: 

172 _PASS_HOLDERS.clear() 

173 

174 

175def clear_loaded_pass_modules() -> None: 

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

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

178 for key in keys_to_remove: 

179 del sys.modules[key]