Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / codegen / inline / strategy_pass.py: 65%

1# Copyright 2026 Huawei Technologies Co., Ltd 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# http://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14# ============================================================================ 

15"""Second inline pass: insert parallel strategies into replaced modules.""" 

16 

17from __future__ import annotations 

18 

19from fnmatch import fnmatch 

20from typing import Collection, Mapping 

21 

22from hyper_parallel.codegen.inline.ir import ( 

23 ForwardExtractPatch, 

24 InlinePatchSet, 

25 SourceObservation, 

26 StrategySpec, 

27) 

28from hyper_parallel.codegen.inline.framework_spec import strategy_spec_for as strategy_spec 

29 

30 

31def rule_block_classes( 

32 rule: SourceObservation, boundary_classes: Mapping[str, str] | None 

33) -> set[str]: 

34 """The concrete classes a rule's match patterns select in the boundary map. 

35 

36 Per-model fact, derived (never registered): the frozen boundary map names 

37 each matched FQN's class, so any custom model class picked up by the 

38 pattern materializes without a framework-side registration. 

39 """ 

40 classes: set[str] = set() 

41 for fqn, class_name in (boundary_classes or {}).items(): 

42 if any(fnmatch(fqn, pattern) for pattern in rule.match): 

43 classes.add(class_name) 

44 return classes 

45 

46 

47def build_strategy_patches( 

48 rules: tuple[SourceObservation, ...], 

49 model_type: str | None = None, 

50 boundary_classes: Mapping[str, str] | None = None, 

51 external_state_classes: Collection[str] | None = None, 

52) -> InlinePatchSet: 

53 """Build inline forward patches from strategy targets. 

54 

55 A plan carries one rule per matched FQN, so the same strategy target (and 

56 therefore the same imports and module snippets) is seen once per module it 

57 applies to. Those contributions are deduplicated here — mirroring 

58 ``build_replacement_patches`` — so a module-level snippet such as the 

59 parallel-state accessor is emitted once per generated file instead of once 

60 per matched layer. 

61 

62 Host derivation is deliberately *not* deduplicated per target: rules that 

63 share one strategy target may select different host classes (one EP 

64 compute target serving both dense-MLP and MoE layers), so every rule 

65 derives its hosts and the per-class ``emitted`` bookkeeping keeps one 

66 materialization per host — mirroring how the freeze side derives 

67 ``external_state_classes`` per record. 

68 

69 A spec whose body inlines a forward but carries no literal target class 

70 (the EP strategy) applies to every class the rule's match selects in 

71 ``boundary_classes``: the block class is a per-model fact, so the patch is 

72 materialized per artifact rather than cached in the capability spec. A 

73 spec that carries a ``subclass_shell`` (inline CP) instead materializes an 

74 in-artifact subclass of the attention host plus a scoped constructor swap, 

75 and never a forward extract — the host keeps its compiled boundary. A 

76 declared strategy that selects no boundary fails fast — a silently 

77 unmaterialized strategy is an artifact that only looks distributed. 

78 """ 

79 

80 patch_set = InlinePatchSet() 

81 emitted: dict[tuple[str, str], str] = {} 

82 imported: set[tuple[str, bool]] = set() 

83 subclassed: set[str] = set() 

84 for rule in rules: 

85 for target in (rule.local_compute_target, rule.inner_wrapper_target): 

86 if target is None: 

87 continue 

88 inner_wrapper = target == rule.inner_wrapper_target 

89 spec = strategy_spec(target, model_type, inner_wrapper=inner_wrapper) 

90 if spec is None: 

91 continue 

92 if (target, inner_wrapper) not in imported: 

93 imported.add((target, inner_wrapper)) 

94 patch_set.imports.extend(spec.imports) 

95 patch_set.module_snippets.extend(spec.snippets) 

96 if spec.subclass_shell is not None: 

97 _materialize_cp_forward( 

98 patch_set, spec, rule, target, model_type, 

99 boundary_classes, subclassed, 

100 ) 

101 continue 

102 if spec.body_template is None: 

103 continue 

104 if spec.target_class is not None: 

105 host_classes = {spec.target_class} 

106 else: 

107 host_classes = rule_block_classes(rule, boundary_classes) 

108 if not host_classes: 

109 # Re-emission of an already-normalized meta: the boundary 

110 # entries the first pass stripped are gone, but the pass 

111 # persisted the derived classes in 

112 # ``meta.external_state_classes`` — the surviving record of 

113 # which classes an inline body replaced. (Today the EP 

114 # shell is the only body-bearing strategy, so the whole 

115 # record is this rule's host set; a second body-bearing 

116 # strategy would need a kind filter here.) 

117 host_classes = set(external_state_classes or ()) 

118 if not host_classes: 

119 raise _unmaterialized_strategy_error( 

120 spec.kind, target, rule, model_type 

121 ) 

122 for host_class in sorted(host_classes): 

123 key = (host_class, spec.method_name) 

124 if key in emitted: 

125 if emitted[key] != spec.body_template: 

126 raise ValueError(f"Conflicting inline strategies for {host_class}.{spec.method_name}") 

127 continue 

128 emitted[key] = spec.body_template 

129 patch_set.forward_extracts.append( 

130 ForwardExtractPatch( 

131 class_name=host_class, 

132 method_name=spec.method_name, 

133 body=spec.body_template, 

134 ) 

135 ) 

136 return patch_set 

137 

138 

139def _materialize_cp_forward( 

140 patch_set: InlinePatchSet, 

141 spec: StrategySpec, 

142 rule: SourceObservation, 

143 target: str, 

144 model_type: str | None, 

145 boundary_classes: Mapping[str, str] | None, 

146 subclassed: set[str], 

147) -> None: 

148 """Materialize CP dispatch in the original host class exactly once.""" 

149 host_classes = rule_block_classes(rule, boundary_classes) 

150 if not host_classes: 

151 raise _unmaterialized_strategy_error(spec.kind, target, rule, model_type) 

152 helper_names = [name for patch in spec.imports for name in patch.names] 

153 if len(helper_names) > 1: 

154 raise ValueError( 

155 f"codegen: CP strategy {target!r} carries multiple helper " 

156 f"imports for one subclass shell: {helper_names}" 

157 ) 

158 for host_class in sorted(host_classes): 

159 if host_class in subclassed: 

160 continue 

161 subclassed.add(host_class) 

162 patch_set.forward_extracts.append( 

163 ForwardExtractPatch( 

164 class_name=host_class, 

165 method_name=spec.method_name, 

166 body=spec.subclass_shell, 

167 merge_original_body=True, 

168 ) 

169 ) 

170 

171 

172def _unmaterialized_strategy_error( 

173 kind: str, target: str, rule: SourceObservation, model_type: str | None 

174) -> RuntimeError: 

175 """The error for a declared strategy that selects no boundary module.""" 

176 return RuntimeError( 

177 f"codegen: {kind} strategy {target!r} matched no boundary " 

178 f"module (patterns {list(rule.match)}, model_type=" 

179 f"{model_type!r}) — the declared strategy cannot be " 

180 "materialized; check the match pattern or the plan's " 

181 "boundary coverage" 

182 )