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

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 InlineRule, 

28 ModuleSnippetPatch, 

29) 

30from hyper_parallel.codegen.inline.specs import replacement_spec 

31 

32logger = logging.getLogger(__name__) 

33 

34 

35def build_replacement_patches(rules: tuple[InlineRule, ...], model_type: str | None = None) -> InlinePatchSet: 

36 """Build constructor replacement patches from ``replace_module`` targets. 

37 

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

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

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

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

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

43 

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

45 ``replace_module`` target must be merged: taking only the first would scope 

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

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

48 original class. 

49 

50 Args: 

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

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

53 

54 Returns: 

55 The constructor replacement patches for the matched rules. 

56 """ 

57 

58 patch_set = InlinePatchSet() 

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

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

61 for rule in rules: 

62 target = rule.replace_target 

63 if target is None: 

64 continue 

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

66 if spec is None: 

67 continue 

68 specs_by_target.setdefault(target, spec) 

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

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

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

72 if not scope_attrs: 

73 logger.warning( 

74 "codegen: replace_module %r matched no module and is not expanded", 

75 target, 

76 ) 

77 continue 

78 patch_set.imports.extend(spec.imports) 

79 if spec.replacement_note: 

80 patch_set.module_snippets.append( 

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

82 ) 

83 patch_set.module_snippets.extend(spec.snippets) 

84 patch_set.constructor_replaces.append( 

85 ConstructorReplacePatch( 

86 old_ctor=spec.old_ctor, 

87 new_ctor=spec.new_ctor, 

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

89 keyword_args=spec.keyword_args, 

90 scope_attrs=scope_attrs, 

91 ) 

92 ) 

93 if spec.remove_class: 

94 patch_set.class_removals.append( 

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

96 ) 

97 return patch_set 

98 

99 

100def _scope_attrs(rule: InlineRule) -> tuple[str, ...]: 

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

102 

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

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

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

106 FQN, so the whole FQN is the attribute. 

107 """ 

108 

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