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

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"""Expand real orchestration source, stopping at framework and kernel APIs.""" 

16 

17import ast 

18from dataclasses import dataclass 

19import inspect 

20import textwrap 

21from types import ModuleType 

22from typing import Any, Callable 

23 

24 

25_PRIMITIVE_PREFIXES = ( 

26 "hyper_parallel.core", 

27 "hyper_parallel.collectives", 

28 "hyper_parallel.components.functional", 

29 "hyper_parallel.components.modules", 

30 "hyper_parallel.distributed.expert_parallel.experts", 

31 "hyper_parallel.distributed.context_parallel.collectives", 

32 "hyper_parallel.codegen.runtime", 

33) 

34 

35 

36@dataclass(frozen=True) 

37class ExpandedSource: 

38 """Source and imports required by one expanded dependency closure.""" 

39 

40 imports: tuple[str, ...] 

41 source: str 

42 

43 

44def expand_function(function: Callable[..., Any]) -> ExpandedSource: 

45 """Copy an orchestration function and recursively resolve its globals. 

46 

47 Constants use their source initializer, so generated cache state never 

48 captures live device tensors from the generation process. Unsupported 

49 globals fail at generation rather than becoming unresolved artifact names. 

50 """ 

51 expansion = _Expansion() 

52 expansion.add_function(function, function.__name__) 

53 return ExpandedSource(tuple(sorted(expansion.imports)), "\n\n".join(expansion.definitions)) 

54 

55 

56class _Expansion: 

57 """Resolve one closure with deterministic order and collision checking.""" 

58 

59 def __init__(self) -> None: 

60 self.imports: set[str] = set() 

61 self.definitions: list[str] = [] 

62 self.bindings: dict[str, Any] = {} 

63 

64 def _claim(self, name: str, value: Any) -> bool: 

65 if name in self.bindings: 

66 if self.bindings[name] is not value: 

67 raise ValueError(f"codegen source expansion: conflicting global {name!r}") 

68 return False 

69 self.bindings[name] = value 

70 return True 

71 

72 def add_function(self, function: Callable[..., Any], name: str) -> None: 

73 """Resolve dependencies before emitting a function definition.""" 

74 if not self._claim(name, function): 

75 return 

76 source = textwrap.dedent(inspect.getsource(function)) 

77 tree = ast.parse(source) 

78 node = tree.body[0] 

79 if not isinstance(node, ast.FunctionDef) or node.decorator_list: 

80 raise ValueError(f"codegen source expansion: unsupported decorated function {name!r}") 

81 global_names = {item.id for item in ast.walk(node) if isinstance(item, ast.Name) 

82 and isinstance(item.ctx, ast.Load)} 

83 local_names = set(function.__code__.co_varnames) 

84 for global_name in sorted(global_names - local_names): 

85 if global_name in function.__globals__: 

86 self.add_global(global_name, function.__globals__[global_name], inspect.getmodule(function)) 

87 node.name = name 

88 self.definitions.append(ast.unparse(node)) 

89 

90 def add_global(self, name: str, value: Any, owner: ModuleType) -> None: 

91 """Emit an import, dependency function, or literal source initializer.""" 

92 if inspect.isfunction(value) and _is_orchestration(value): 

93 self.add_function(value, name) 

94 return 

95 if not self._claim(name, value): 

96 return 

97 if isinstance(value, ModuleType): 

98 self.imports.add(f"import {value.__name__} as {name}") 

99 return 

100 module_name = getattr(value, "__module__", None) 

101 original_name = getattr(value, "__name__", None) 

102 if module_name and original_name: 

103 self.imports.add(f"from {module_name} import {original_name} as {name}") 

104 return 

105 initializer = _literal_initializer(owner, name) 

106 self.definitions.append(f"{name} = {initializer}") 

107 

108 

109def _is_orchestration(function: Callable[..., Any]) -> bool: 

110 module = function.__module__ 

111 return module.startswith("hyper_parallel.") and not any( 

112 module == prefix or module.startswith(prefix + ".") for prefix in _PRIMITIVE_PREFIXES 

113 ) 

114 

115 

116def _literal_initializer(owner: ModuleType, name: str) -> str: 

117 """Return the source text of the literal assigned to ``name`` in ``owner``.""" 

118 for node in ast.parse(inspect.getsource(owner)).body: 

119 if isinstance(node, ast.Assign): 

120 targets, value = node.targets, node.value 

121 elif isinstance(node, ast.AnnAssign): 

122 targets, value = [node.target], node.value 

123 else: 

124 continue 

125 if any(isinstance(target, ast.Name) and target.id == name for target in targets): 

126 try: 

127 ast.literal_eval(value) 

128 except (ValueError, TypeError) as exc: 

129 raise ValueError(f"codegen source expansion: nonliteral global {name!r}") from exc 

130 return ast.unparse(value) 

131 raise ValueError(f"codegen source expansion: no literal initializer for {name!r}")