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

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"""Second inline pass: insert parallel strategies into replaced modules.""" 

16 

17from __future__ import annotations 

18 

19from hyper_parallel.codegen.inline.ir import ForwardExtractPatch, InlinePatchSet, InlineRule 

20from hyper_parallel.codegen.inline.specs import strategy_spec 

21 

22 

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

24 """Build inline forward patches from strategy targets. 

25 

26 A plan carries one rule per matched FQN, so the same strategy target (and 

27 therefore the same imports and module snippets) is seen once per module it 

28 applies to. Those contributions are deduplicated here — mirroring 

29 ``build_replacement_patches`` — so a module-level snippet such as the 

30 parallel-state accessor is emitted once per generated file instead of once 

31 per matched layer. 

32 """ 

33 

34 patch_set = InlinePatchSet() 

35 emitted: dict[tuple[str, str], str] = {} 

36 seen: set[tuple[str, bool]] = set() 

37 for rule in rules: 

38 for target in (rule.local_compute_target, rule.inner_wrapper_target): 

39 if target is None: 

40 continue 

41 inner_wrapper = target == rule.inner_wrapper_target 

42 if (target, inner_wrapper) in seen: 

43 continue 

44 spec = strategy_spec(target, model_type, inner_wrapper=inner_wrapper) 

45 if spec is None: 

46 continue 

47 seen.add((target, inner_wrapper)) 

48 patch_set.imports.extend(spec.imports) 

49 patch_set.module_snippets.extend(spec.snippets) 

50 if spec.target_class is None: 

51 continue 

52 key = (spec.target_class, spec.method_name) 

53 if key in emitted: 

54 if emitted[key] != spec.body_template: 

55 raise ValueError(f"Conflicting inline strategies for {spec.target_class}.{spec.method_name}") 

56 continue 

57 emitted[key] = spec.body_template 

58 patch_set.forward_extracts.append( 

59 ForwardExtractPatch( 

60 class_name=spec.target_class, 

61 method_name=spec.method_name, 

62 body=spec.body_template, 

63 ) 

64 ) 

65 return patch_set