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

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"""First inline pass: sink YAML module replacements into source constructors.""" 

16 

17from __future__ import annotations 

18 

19import logging 

20from collections import defaultdict 

21from typing import Any 

22 

23from hyper_parallel.codegen.inline.ir import ( 

24 ClassRemovalPatch, 

25 ConstructorReplacePatch, 

26 InlinePatchSet, 

27 SourceObservation, 

28 ModuleSnippetPatch, 

29) 

30from hyper_parallel.codegen.inline.framework_spec import replacement_spec_for as replacement_spec 

31 

32logger = logging.getLogger(__name__) 

33 

34 

35def build_replacement_patches( 

36 rules: tuple[SourceObservation, ...], model_type: str | None = None 

37) -> InlinePatchSet: 

38 """Build constructor replacement patches from module actions. 

39 

40 The rule's settled FQNs (``rule.match``, recorded from the same plan the 

41 native path applies) define the *scope* of each rewrite: only the 

42 ``self.<attr>`` call sites those FQNs name are sunk. A rule that matched no 

43 module is not expanded at all -- generating a replacement the plan never 

44 decided would make the artifact diverge from the trained model. 

45 

46 The plan records one entry per matched FQN, so the entries sharing a 

47 module implementation must be merged: taking only the first would scope 

48 the rewrite to a single matched attribute (``input_layernorm``) and silently 

49 leave its siblings (``post_attention_layernorm``, ``model.norm``) on the 

50 original class. 

51 

52 Args: 

53 rules: Inline rules recovered from the plan's frozen metadata. 

54 model_type: Model family key used to resolve framework declarations. 

55 

56 Returns: 

57 The constructor replacement patches for the matched rules. 

58 """ 

59 

60 patch_set = InlinePatchSet() 

61 scopes_by_target: dict[str, set[str]] = defaultdict(set) 

62 fqns_by_target: dict[str, set[str]] = defaultdict(set) 

63 specs_by_target: dict[str, Any] = {} 

64 for rule in rules: 

65 target = rule.replace_target 

66 if target is None: 

67 continue 

68 spec = replacement_spec(target, model_type, module_type=rule.module_type) 

69 if spec is None: 

70 continue 

71 specs_by_target.setdefault(target, spec) 

72 scopes_by_target[target].update(_scope_attrs(rule)) 

73 fqns_by_target[target].update(fqn for fqn in rule.match if fqn) 

74 for target, spec in specs_by_target.items(): 

75 scope_attrs = tuple(sorted(scopes_by_target[target])) 

76 if not scope_attrs: 

77 logger.warning( 

78 "codegen: module replacement %r matched no module and is not expanded", 

79 target, 

80 ) 

81 continue 

82 patch_set.imports.extend(spec.imports) 

83 if spec.replacement_note: 

84 patch_set.module_snippets.append( 

85 ModuleSnippetPatch(f"# Codegen replacement: {spec.replacement_note}") 

86 ) 

87 patch_set.module_snippets.extend(spec.snippets) 

88 patch_set.constructor_replaces.append( 

89 ConstructorReplacePatch( 

90 old_ctor=spec.old_ctor, 

91 new_ctor=spec.new_ctor, 

92 mode=spec.mode, # type: ignore[arg-type] 

93 keyword_args=spec.keyword_args, 

94 scope_attrs=scope_attrs, 

95 scope_fqns=tuple(sorted(fqns_by_target[target])), 

96 ) 

97 ) 

98 if spec.remove_class: 

99 patch_set.class_removals.append( 

100 ClassRemovalPatch(old_name=spec.old_ctor, new_name=spec.new_ctor) 

101 ) 

102 return patch_set 

103 

104 

105def _scope_attrs(rule: SourceObservation) -> tuple[str, ...]: 

106 """Attribute names of the rule's settled FQNs, i.e. the rewrite scope. 

107 

108 ``model.layers.3.post_attention_layernorm`` scopes the rewrite to 

109 ``self.post_attention_layernorm = <ctor>(...)`` assignments, which is the 

110 granularity native's FQN matching produces. ``model.norm`` is a sole-name 

111 FQN, so the whole FQN is the attribute. 

112 """ 

113 

114 return tuple(sorted({fqn.rsplit(".", 1)[-1] for fqn in rule.match if fqn}))