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

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"""Apply source-visible forward extraction patches.""" 

16 

17from __future__ import annotations 

18 

19import textwrap 

20from collections.abc import Iterable 

21 

22from hyper_parallel.codegen.astkit.edits import TextEdit, replace_function_body 

23from hyper_parallel.codegen.astkit.index import FunctionInfo, build_source_index 

24from hyper_parallel.codegen.emit.parallel import _build_forward_impl_edit 

25from hyper_parallel.codegen.inline.ir import ForwardExtractPatch 

26 

27_HIDDEN_ARG_TOKEN = "__hp_hidden_arg__" 

28 

29 

30def forward_extract_edits( 

31 source_text: str, 

32 patches: Iterable[ForwardExtractPatch], 

33) -> list[TextEdit]: 

34 """Build edits that preserve original forwards and install dispatch bodies.""" 

35 patch_list = tuple(patches) 

36 if not patch_list: 

37 return [] 

38 index = build_source_index(source_text) 

39 edits: list[TextEdit] = [] 

40 for patch in patch_list: 

41 cls = index.find_class(patch.class_name) 

42 if cls is None: 

43 continue 

44 func = cls.methods.get(patch.method_name) 

45 if func is None: 

46 continue 

47 body = _bind_forward_args(_bind_hidden_arg(patch.body, func), func) 

48 prefix = _parallel_state_prefix(body) 

49 if patch.merge_original_body: 

50 edits.extend(_merged_forward_edits(source_text, patch, func, body, prefix)) 

51 continue 

52 impl = _build_forward_impl_edit(source_text, func, prefix=prefix) 

53 if impl is not None: 

54 edits.append(impl) 

55 edits.extend(replace_function_body(source_text, func, body, indent=" ")) 

56 return edits 

57 

58 

59def _parallel_state_prefix(body: str) -> str: 

60 """Render static enable flags read by a generated dispatch body.""" 

61 lines = [ 

62 f" {axis}_enable = False" 

63 for axis in ("tp", "cp", "ep") 

64 if f"self.{axis}_enable" in body 

65 ] 

66 return "\n".join(lines) + "\n\n" if lines else "" 

67 

68 

69def _merged_forward_edits( 

70 source_text: str, 

71 patch: ForwardExtractPatch, 

72 func: FunctionInfo, 

73 body: str, 

74 prefix: str, 

75) -> list[TextEdit]: 

76 """Merge dispatch with the original body when no separate impl is needed.""" 

77 fallback = _forward_impl_fallback(func) 

78 if fallback not in body: 

79 raise RuntimeError( 

80 f"codegen: merged strategy for {patch.class_name}." 

81 f"{patch.method_name} has no original-forward fallback" 

82 ) 

83 dispatch = body.replace(fallback, "", 1).rstrip() 

84 original = textwrap.dedent(source_text[func.body_start:func.body_end]).rstrip() 

85 edits = list(replace_function_body( 

86 source_text, func, f"{dispatch}\n\n{original}", indent=" " 

87 )) 

88 if prefix: 

89 edits.insert(0, TextEdit(func.def_offset, func.def_offset, prefix)) 

90 return edits 

91 

92 

93def _bind_hidden_arg(body: str, func: FunctionInfo) -> str: 

94 """Bind a strategy shell's hidden-input placeholder to its host signature.""" 

95 if _HIDDEN_ARG_TOKEN not in body: 

96 return body 

97 if not func.param_names: 

98 raise RuntimeError( 

99 f"codegen: inline strategy body references {_HIDDEN_ARG_TOKEN!r} " 

100 "but the target method has no positional parameter to bind it to" 

101 ) 

102 return body.replace(_HIDDEN_ARG_TOKEN, func.param_names[0]) 

103 

104 

105def _bind_forward_args(body: str, func: FunctionInfo) -> str: 

106 """Bind a dispatch shell's argument token to the original signature.""" 

107 token = "__hp_forward_args__" 

108 if token not in body: 

109 return body 

110 args = list(func.param_names) 

111 if func.vararg_name: 

112 args.append(f"*{func.vararg_name}") 

113 args.extend(f"{name}={name}" for name in func.kwonly_names) 

114 if func.kwarg_name: 

115 args.append(f"**{func.kwarg_name}") 

116 return body.replace(token, ", ".join(args)) 

117 

118 

119def _forward_impl_fallback(func: FunctionInfo) -> str: 

120 """Render the dispatch-shell fallback that calls the extracted method.""" 

121 args = list(func.param_names) 

122 if func.vararg_name: 

123 args.append(f"*{func.vararg_name}") 

124 args.extend(f"{name}={name}" for name in func.kwonly_names) 

125 if func.kwarg_name: 

126 args.append(f"**{func.kwarg_name}") 

127 return f"return self._forward_impl({', '.join(args)})" 

128 

129 

130__all__ = ["forward_extract_edits"]