Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / compile / passes / overlap / utils.py: 97%

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"""Operator and effect classification shared by the overlap schedulers.""" 

16 

17import operator 

18 

19from torch import fx 

20from torch.ops import _c10d_functional, aten 

21 

22 

23def is_op(node: fx.Node, packet: object) -> bool: 

24 """Match both functional collective packets and their default overloads.""" 

25 return node.op == "call_function" and node.target in (packet, packet.default) 

26 

27 

28def is_view(node: fx.Node) -> bool: 

29 """Recognize immutable aliases, including tuple projections of split views.""" 

30 if node.op != "call_function": 

31 return False 

32 if node.target is operator.getitem: 

33 return True 

34 schema = getattr(node.target, "_schema", None) 

35 return schema is not None and not schema.is_mutable and any( 

36 result.alias_info is not None for result in schema.returns 

37 ) 

38 

39 

40def is_collective(node: fx.Node) -> bool: 

41 """Recognize functional collective launches, excluding their waits.""" 

42 return node.op == "call_function" and ( 

43 str(node.target).startswith(("_c10d_functional.", "_c10d_functional_autograd.")) 

44 and not is_op(node, _c10d_functional.wait_tensor) 

45 ) 

46 

47 

48def is_barrier(node: fx.Node) -> bool: 

49 """Bound reordering at mutations and unknown effects. 

50 

51 Tagged FSDP storage releases are handled by the owning scheduler's data 

52 dependencies; unrelated parameter releases need not serialize communication. 

53 """ 

54 if node.op in ("placeholder", "get_attr", "output") or is_collective(node): 

55 return False 

56 if is_op(node, _c10d_functional.wait_tensor): 

57 return False 

58 if node.meta.get("comm_type") in ("fsdp_reshard_free", "fsdp_reshard_free_backward"): 

59 return False 

60 if node.op != "call_function": 

61 return True 

62 schema = getattr(node.target, "_schema", None) 

63 if schema is not None: 

64 return schema.is_mutable 

65 return node.target is not operator.getitem 

66 

67 

68def rs_layout_preparation(launch: fx.Node) -> list[fx.Node]: 

69 """Recognize a private default contiguous input owned by one backward RS. 

70 

71 This read-only predicate is shared by metadata, planning, rewrite and overlap. 

72 Shared, explicitly forward, non-default or differently owned nodes stay put. 

73 

74 Args: 

75 launch: Reduce-scatter launch whose input preparation is inspected. 

76 """ 

77 inputs = launch.all_input_nodes 

78 if len(inputs) != 1: 

79 return [] 

80 source = inputs[0] 

81 if set(source.users) != {launch} or source.meta.get("autograd_backward") is False: 

82 return [] 

83 owner = source.meta.get("fsdp_reduce_scatter_prepare") 

84 if owner is not None and owner != launch.name: 

85 return [] 

86 if len(source.args) != 1 or not isinstance(source.args[0], fx.Node) or source.kwargs: 

87 return [] 

88 if source.op == "call_method" and source.target == "contiguous": 

89 return [source] 

90 if is_op(source, aten.contiguous) and owner == launch.name: 

91 return [source] 

92 return [] 

93 

94 

95def bucket_preparation(launch: fx.Node, order: dict[fx.Node, int]) -> list[fx.Node]: 

96 """Move only the private, pure pack DAG emitted by the AG/RS bucket transforms. 

97 

98 Args: 

99 launch: Bucket collective owning the pack nodes. 

100 order: Current graph node positions for topological ordering. 

101 """ 

102 bucket_id = launch.meta.get("fsdp_bucket_id") 

103 if not bucket_id: 

104 return [] 

105 preparation = set() 

106 stack = list(launch.all_input_nodes) 

107 while stack: 

108 node = stack.pop() 

109 if node.meta.get("fsdp_bucket_prepare") != bucket_id: 

110 continue 

111 schema = getattr(node.target, "_schema", None) 

112 if schema is None or schema.name not in ("aten::reshape", "aten::cat", "aten::contiguous"): 

113 return [] 

114 if node not in preparation: 

115 preparation.add(node) 

116 stack.extend(node.all_input_nodes) 

117 if any(set(node.users) - preparation - {launch} for node in preparation): 

118 return [] 

119 return sorted(preparation, key=order.get)