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

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"""Intermediate representation for inline generated-model patches.""" 

16 

17from __future__ import annotations 

18 

19from dataclasses import dataclass, field 

20from typing import Literal, Optional 

21 

22 

23@dataclass(frozen=True) 

24class SourceObservation: 

25 """One normalized replacement or injection observed during native build.""" 

26 

27 match: tuple[str, ...] 

28 module_type: Optional[str] = None 

29 replace_target: Optional[str] = None 

30 inner_wrapper_target: Optional[str] = None 

31 local_compute_target: Optional[str] = None 

32 

33 

34@dataclass(frozen=True) 

35class ImportPatch: 

36 """A ``from module import names`` insertion.""" 

37 

38 module: str 

39 names: tuple[str, ...] 

40 raw: str | None = None 

41 

42 

43@dataclass(frozen=True) 

44class ModuleSnippetPatch: 

45 """A top-level source snippet inserted after imports.""" 

46 

47 text: str 

48 

49 

50@dataclass(frozen=True) 

51class ConstructorReplacePatch: 

52 """Replace constructor calls in source text. 

53 

54 ``scope_attrs`` carries the attribute names of the rule's *matched* FQNs 

55 (``{"input_layernorm", "post_attention_layernorm", "norm"}`` for the 

56 RMSNorm rule). Native replaces module instances by FQN, so a whole-file 

57 class-name rewrite would silently replace call sites the rule never 

58 matched -- for example the Q/K norm inside the attention class that the 

59 recipe deliberately leaves as the original implementation. An empty scope 

60 means the caller imposes no restriction; the YAML-driven pass never emits 

61 an unscoped patch. 

62 """ 

63 

64 old_ctor: str 

65 new_ctor: str 

66 mode: Literal["name", "wrap_source"] = "name" 

67 keyword_args: tuple[str, ...] = () 

68 scope_attrs: tuple[str, ...] = () 

69 scope_fqns: tuple[str, ...] = () 

70 

71 

72@dataclass(frozen=True) 

73class ForwardExtractPatch: 

74 """Keep the original method as ``_forward_impl`` and replace its body.""" 

75 

76 class_name: str 

77 method_name: str 

78 body: str 

79 merge_original_body: bool = False 

80 

81 

82@dataclass(frozen=True) 

83class ClassRemovalPatch: 

84 """Remove an unused original HF class definition and rewrite its references. 

85 

86 ``old_name`` is the class being removed; ``new_name`` is the fused 

87 replacement that any surviving bare reference (e.g. inside 

88 ``_can_record_outputs`` or ``isinstance`` checks) should be rewritten to. 

89 Constructor calls are handled separately by ``ConstructorReplacePatch``. 

90 """ 

91 

92 old_name: str 

93 new_name: str 

94 

95 

96@dataclass 

97class InlinePatchSet: 

98 """All edits emitted by one inline pass.""" 

99 

100 imports: list[ImportPatch] = field(default_factory=list) 

101 module_snippets: list[ModuleSnippetPatch] = field(default_factory=list) 

102 constructor_replaces: list[ConstructorReplacePatch] = field(default_factory=list) 

103 forward_extracts: list[ForwardExtractPatch] = field(default_factory=list) 

104 class_removals: list[ClassRemovalPatch] = field(default_factory=list) 

105 

106 

107@dataclass(frozen=True) 

108class ReplacementSpec: 

109 """Source-level form of one module replacement implementation.""" 

110 

111 old_ctor: str 

112 new_ctor: str 

113 imports: tuple[ImportPatch, ...] 

114 mode: str = "name" 

115 keyword_args: tuple[str, ...] = () 

116 snippets: tuple[ModuleSnippetPatch, ...] = () 

117 remove_class: bool = True 

118 replacement_note: str | None = None 

119 

120 

121@dataclass(frozen=True) 

122class StrategySpec: 

123 """Source-level form of one parallel strategy target.""" 

124 

125 kind: str 

126 imports: tuple[ImportPatch, ...] 

127 target_class: str | None = None 

128 method_name: str = "forward" 

129 body_template: str | None = None 

130 subclass_shell: str | None = None 

131 subclass_state_attr: str | None = None 

132 snippets: tuple[ModuleSnippetPatch, ...] = () 

133 strip_boundary_subpatterns: tuple[str, ...] = () 

134 

135 def __post_init__(self) -> None: 

136 """Require a complete method patch or an imports-only declaration.""" 

137 if self.target_class is not None and self.body_template is None: 

138 raise ValueError("Inline strategy target_class requires a body_template") 

139 if self.target_class is not None and (not self.target_class or not self.body_template.strip()): 

140 raise ValueError("Inline strategy class and body must be nonempty") 

141 if not self.method_name.isidentifier(): 

142 raise ValueError("Inline strategy method_name must be a Python identifier") 

143 if self.subclass_shell is not None: 

144 if self.target_class is not None or self.body_template is not None: 

145 raise ValueError( 

146 "Inline strategy subclass_shell is exclusive with " 

147 "target_class and body_template" 

148 ) 

149 if not self.subclass_shell.strip(): 

150 raise ValueError("Inline strategy subclass shell must be nonempty")