Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_multicore / scheduler / graph.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""" 

16Compute graph abstractions: OpType, TensorSpec, SplitSpec, OperatorNode, ComputeGraph. 

17 

18ComputeGraph.propagate_splits() traverses the DAG in topological order and 

19automatically computes task_num / split_dim for every node from their SplitSpec. 

20""" 

21from collections import deque 

22from dataclasses import dataclass, field 

23from enum import Enum 

24from typing import Any, Callable, List, Optional, Tuple, Union 

25 

26from hyper_multicore.scheduler.config import TaskSplitValue 

27 

28 

29class OpType(Enum): 

30 """Operator kinds supported by the multicore compute graph.""" 

31 

32 ALLTOALL = "alltoall" 

33 AIV_PIPELINE = "aiv_pipeline" 

34 GMM = "gmm" 

35 SWIGLU = "swiglu" 

36 SWIGLU_GRAD = "swiglu_grad" 

37 MHC_POST = "mhc_post" 

38 MHC_NORM_CAST = "mhc_norm_cast" 

39 MHC_PROJECTION = "mhc_projection" 

40 MHC_MAPPING = "mhc_mapping" 

41 MHC_INPUT_MIX = "mhc_input_mix" 

42 MHC_PIPELINE = "mhc_pipeline" 

43 RMS_NORM = "rms_norm" 

44 

45 

46@dataclass 

47class TensorSpec: 

48 """Describes a single tensor in the graph (a graph edge / shared buffer).""" 

49 name: str 

50 shape: list 

51 param_position: int 

52 dtype_size: int = 2 # bf16=2, int64=8, int32=4 

53 tensor_type: int = 0 # 0=tensor, 1=tensorlist 

54 transpose: bool = False 

55 is_dynamic: bool = False 

56 

57 # Set by propagate_splits(); not user-supplied. 

58 split_dim: int = field(default=-1, init=False) 

59 split_num: int = field(default=1, init=False) 

60 

61 

62@dataclass 

63class SplitSpec: 

64 """ 

65 Declarative split specification attached to an OperatorNode. 

66 

67 split_inputs 

68 None → source operator; always computes via task_num_fn. 

69 list → list of (input_idx, dim) pairs; ALL must already be split on 

70 the given dim for this op to be active. 

71 

72 split_output_dims 

73 Per-output split axis. -1 = leave un-split. 

74 

75 task_num_fn 

76 Callable(tsv) -> int. Returns task_num when the split condition holds. 

77 """ 

78 split_inputs: Optional[List[Tuple[int, int]]] 

79 task_num_fn: Callable 

80 split_output_dims: List[int] = field(default_factory=lambda: [0]) 

81 

82 

83@dataclass 

84class OperatorNode: 

85 """A single operator node in the compute DAG.""" 

86 name: str 

87 op_type: OpType 

88 inputs: List[TensorSpec] 

89 outputs: List[TensorSpec] 

90 param_positions: List[int] 

91 split_value: int 

92 split_spec: SplitSpec 

93 tiling_position: int 

94 fill_config: Any # FillConfig subclass instance 

95 kernel_spec: Any = None # reserved for @MultiCore decorator path 

96 diagnostic_name: Optional[str] = None # optional Host-side diagnostic display name 

97 diagnostic_view_resolver: Optional[Callable[..., Any]] = None 

98 diagnostic_host_ports: bool = False # Native hard-coded ABI operands, not TaskDesc TensorDesc arrays. 

99 

100 predecessors: List['OperatorNode'] = field(default_factory=list) 

101 successors: List['OperatorNode'] = field(default_factory=list) 

102 task_num: int = field(default=0, init=False) 

103 

104 

105class ComputeGraph: 

106 """Directed acyclic graph describing operator execution order.""" 

107 

108 def __init__(self) -> None: 

109 """Initialize an empty compute graph.""" 

110 self._nodes: dict[str, OperatorNode] = {} 

111 self._insertion_order: list[str] = [] 

112 

113 def add_op(self, op: OperatorNode) -> 'ComputeGraph': 

114 """Add an operator node to the graph. 

115 

116 Args: 

117 op: Operator node to append. 

118 """ 

119 self._nodes[op.name] = op 

120 self._insertion_order.append(op.name) 

121 return self 

122 

123 def add_edge( 

124 self, src: Union[OperatorNode, str], dst: Union[OperatorNode, str] 

125 ) -> 'ComputeGraph': 

126 """Connect two operator nodes. 

127 

128 Args: 

129 src: Source node or its registered name. 

130 dst: Destination node or its registered name. 

131 """ 

132 s = src if isinstance(src, OperatorNode) else self._nodes[src] 

133 d = dst if isinstance(dst, OperatorNode) else self._nodes[dst] 

134 s.successors.append(d) 

135 d.predecessors.append(s) 

136 return self 

137 

138 def get_op(self, name: str) -> OperatorNode: 

139 """Retrieve an operator node by name. 

140 

141 Args: 

142 name: Registered operator name. 

143 """ 

144 return self._nodes[name] 

145 

146 def topological_sort(self) -> List[OperatorNode]: 

147 """Kahn's algorithm; insertion order used for tie-breaking.""" 

148 in_deg = {n: len(op.predecessors) for n, op in self._nodes.items()} 

149 queue = deque(n for n in self._insertion_order if in_deg[n] == 0) 

150 order = [] 

151 while queue: 

152 n = queue.popleft() 

153 order.append(self._nodes[n]) 

154 for succ in self._nodes[n].successors: 

155 in_deg[succ.name] -= 1 

156 if in_deg[succ.name] == 0: 

157 queue.append(succ.name) 

158 if len(order) != len(self._nodes): 

159 raise ValueError("ComputeGraph has a cycle") 

160 return order 

161 

162 def propagate_splits(self, tsv: TaskSplitValue) -> None: 

163 """ 

164 Compute task_num and split_dim for every node from their SplitSpec. 

165 

166 Must be called once after graph construction and before any fill loop. 

167 Resets all TensorSpec split annotations before traversal. 

168 

169 Args: 

170 tsv: Topology values used to resolve split counts. 

171 """ 

172 for op in self._nodes.values(): 

173 for t in op.inputs + op.outputs: 

174 t.split_dim = -1 

175 t.split_num = 1 

176 

177 for op in self.topological_sort(): 

178 ss = op.split_spec 

179 if ss.split_inputs is None: 

180 task_num = ss.task_num_fn(tsv) 

181 else: 

182 if all(op.inputs[idx].split_dim == dim for idx, dim in ss.split_inputs): 

183 task_num = ss.task_num_fn(tsv) 

184 else: 

185 task_num = 1 

186 

187 op.task_num = task_num 

188 

189 for i, out in enumerate(op.outputs): 

190 if i < len(ss.split_output_dims): 

191 d = ss.split_output_dims[i] 

192 out.split_dim = d if (task_num > 1 and d >= 0) else -1 

193 else: 

194 out.split_dim = -1 

195 out.split_num = task_num