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

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

25 """One normalized rule recovered from frozen YAML-derived metadata.""" 

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 

70 

71@dataclass(frozen=True) 

72class ForwardExtractPatch: 

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

74 

75 class_name: str 

76 method_name: str 

77 body: str 

78 

79 

80@dataclass(frozen=True) 

81class ClassRemovalPatch: 

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

83 

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

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

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

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

88 """ 

89 

90 old_name: str 

91 new_name: str 

92 

93 

94@dataclass 

95class InlinePatchSet: 

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

97 

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

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

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

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

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