Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / sapp_ppb / simulator / sim_block.py: 95%
225 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-07-06 05:41 +0800
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"""Simulator block primitives: compute blocks, comm (send/receive) blocks, head sentinels."""
16from __future__ import annotations
18from dataclasses import dataclass, field
19from typing import List, Tuple
21import matplotlib.pyplot as plt
22import numpy as np
23from matplotlib.patches import Polygon, Rectangle
25from hyper_parallel.auto_parallel.sapp_ppb.simulator.utils import color_mix, dfs_builder
28@dataclass
29class BlockSim:
30 r"""base block sim class"""
31 stage: int # p
32 state: str # s
33 id: int # m
34 chunk: int # v
35 time: float
36 type: str
37 start: float = None
38 end: float = None
39 pre: BlockSim = field(repr=False, default=None)
40 left: BlockSim = field(repr=False, default=None)
41 # pylint: disable=E0601
42 right: MicroBlockSim = field(repr=False, default=None)
43 depend_pre: BlockSim = field(repr=False, default=None)
44 depend_left: BlockSim = field(repr=False, default=None)
45 finish = False
46 in_queue = False
47 flag = False
48 _color = '0;38'
49 father: BlockSim = field(repr=False, default=None)
51 @property
52 def label(self) -> tuple:
53 """Identity tuple ``(type, state, id, chunk, stage)`` used in log messages."""
54 res = (self.type, self.state, self.id, self.chunk, self.stage)
55 return res
57 @property
58 def color_label(self) -> str:
59 """ANSI-coloured version of :attr:`label` for terminal output."""
60 return f"\033[{self._color}m{self.label}\033[0m"
62 @dfs_builder(False)
63 def build_without_comm(self) -> None:
64 r"""Build pipeline timeline without comm blocks and dependency."""
65 self.pre.build_without_comm()
66 self.left.build_without_comm()
67 self.start = max(self.pre.end, self.left.end)
68 self.end = self.start + self.time
70 @dfs_builder(True)
71 def build_with_comm(self) -> None:
72 r"""Build pipeline timeline with comm blocks and dependency."""
73 self.depend_pre.build_with_comm()
74 self.depend_left.build_with_comm()
75 self.start = max(self.depend_pre.end, self.depend_left.end)
76 self.end = self.start + self.time
78 def reset_time(self) -> None:
79 r"""reset time"""
80 self.start = None
81 self.end = None
82 self.finish = False
84 # pylint: disable=W0613
85 def loc_size(self, x: float = 0, equal_wide: bool = False,
86 mode: str = 'compact') -> Tuple[float, float, float, float]:
87 """Return ``(x, y, dx, dy)`` for plotting this block at x-coordinate ``x``."""
88 x = x if self.start is None else self.start
89 dx = 1 if equal_wide else self.time
90 res = x, self.stage + 0.5, dx, 1
91 return res
93 def loop(self, comm: bool = False) -> List["BlockSim"]:
94 """Return the first dependency cycle discovered via DFS, or an empty list."""
95 if self.flag and not self.in_queue:
96 return []
97 res = []
98 if self.in_queue:
99 loop = [self]
100 block = self.father
101 while block.father and block is not self:
102 block = block.father
103 loop.append(block)
104 return loop
105 self.flag = True
106 self.in_queue = True
107 depends = [self.depend_pre, self.depend_left] if comm else [self.pre, self.left]
108 for dep in depends:
109 if dep:
110 dep.father = self
111 res.extend(dep.loop(comm=comm))
112 dep.father = None
113 self.in_queue = False
114 return res
116 def comm_loop(self) -> list[BlockSim]:
117 r"""recursively check comm loop"""
118 return self.loop(True)
121@dataclass
122class HeadBlockSim(BlockSim):
123 r"""sim block of head"""
124 stage: int # p
125 type: str = 'h'
126 id: int = field(repr=False, init=False)
127 state: str = field(repr=False, init=False)
128 chunk: int = field(repr=False, init=False)
129 time: float = 0.
130 start: float = 0.
131 end: float = 0.
132 finish = True
134 @property
135 def label(self) -> tuple:
136 """Identity tuple ``(type, stage)`` for head sentinel blocks."""
137 return (self.type, self.stage)
139 @property
140 def repr(self) -> str:
141 """Multi-line representation listing every block chained to the right of this head."""
142 s_list = []
143 block = self
144 while block:
145 s_list.append(repr(block))
146 block = block.right
147 return '\n'.join(s_list)
149 # pylint: disable=W0613
150 def draw(self, ax: plt.Axes, *args: object, **kwargs: object) -> None:
151 """No-op: the head sentinel is not rendered."""
152 return
154 def build_without_comm(self) -> None:
155 """No-op: the head sentinel has no non-comm dependencies."""
156 return
158 def build_with_comm(self) -> None:
159 """No-op: the head sentinel has no comm dependencies."""
160 return
162 def reset_time_recursive(self) -> None:
163 """No-op: the head sentinel has no times to reset."""
164 return
167@dataclass
168class MicroBlockSim(BlockSim):
169 r"""compute sim block"""
170 type: str = 'c'
171 mem: float = 0.
172 mem_par: float = 0.
173 phase: str = None
174 # pylint: disable=E0601
175 send_block: SendBlockSim = field(repr=False, default=None)
176 # pylint: disable=E0601
177 rec_block: RecBlockSim = field(repr=False, default=None)
179 def __post_init__(self) -> None:
180 """Choose the ANSI colour based on forward / backward state."""
181 self._color = '1;34' if self.state == 'f' else '1;33'
183 # pylint: disable=W0613
184 def draw(self, ax: plt.Axes, *args: object, **kwargs: object) -> None:
185 """Render this compute block as a coloured rectangle on ``ax``."""
186 x, y, dx, dy = self.loc_size(kwargs.get('index', 0), kwargs.get('equal_wide', False))
187 color = (167 / 255, 184 / 255, 231 / 255) if self.state == 'f' else (255 / 255, 213 / 255, 143 / 255)
188 mix_color = (240 / 255, 255 / 255, 245 / 255) if self.state == 'f' else (255 / 255, 240 / 255, 255 / 255)
189 color = color_mix(mix_color, color, w1=self.chunk / 3)
190 if self.phase == 'warmup' and kwargs.get('phase', False):
191 edgecolor = 'lightblue'
192 elif self.phase == 'cooldown' and kwargs.get('phase', False):
193 edgecolor = 'orange'
194 else:
195 edgecolor = 'black'
196 rect = Rectangle((x, y - dy / 2), dx, dy, facecolor=color, edgecolor=edgecolor, linewidth=0.4)
197 if dx > 0.008 * kwargs.get('width', 0):
198 ax.text(rect.xy[0] + dx / 2, rect.xy[1] + dy / 2, str(self.id),
199 ha='center', va='center', color='black', fontdict={'fontsize': 9})
200 ax.add_patch(rect)
202 def reset_time_recursive(self) -> None:
203 r"""reset block time"""
204 if self.finish:
205 self.pre.reset_time_recursive()
206 self.left.reset_time_recursive()
207 self.reset_time()
210@dataclass
211class CommBlockSim(BlockSim):
212 r"""sim comm block"""
213 host: MicroBlockSim = field(repr=False, default=None)
214 dual: CommBlockSim = field(repr=False, default=None)
216 def get_triangle(self, x: float, y: float, dx: float,
217 dy: float) -> List[List[float]]:
218 """Return the three triangle vertices used by :meth:`draw`."""
219 raise NotImplementedError
221 # pylint: disable=W0613
222 def draw(self, ax: plt.Axes, *args: object, **kwargs: object) -> None:
223 """Render this communication block as a coloured triangle on ``ax``."""
224 color = (167 / 255, 184 / 255, 231 / 255) if self.state == 'f' else (255 / 255, 213 / 255, 143 / 255)
225 mix_color = (240 / 255, 255 / 255, 255 / 255) if self.state == 'f' else (255 / 255, 240 / 255, 255 / 255)
226 color = color_mix(mix_color, color, w1=1.2 * self.chunk / 3)
227 index, equal_wide, mode = (kwargs.get('index', 0), kwargs.get('equal_wide', False),
228 kwargs.get('mode', 'compact'))
229 x, y, dx, dy = self.loc_size(index, equal_wide, mode)
230 xy = self.get_triangle(x, y, dx, dy)
231 tri = Polygon(xy, closed=True, facecolor=color, edgecolor='black', linewidth=0.4)
232 ax.add_patch(tri)
235@dataclass
236class SendBlockSim(CommBlockSim):
237 r"""sim send comm block"""
238 type: str = 's'
239 _color = '35'
241 def loc_size(self, x: float = 0, equal_wide: bool = False,
242 mode: str = 'compact') -> Tuple[float, float, float, float]:
243 """Return ``(x, y, dx, dy)`` for a send block relative to its compute host."""
244 host_x, _, hostdx_, _ = self.host.loc_size(x, equal_wide)
245 x, y, _, _ = super().loc_size(x, equal_wide)
246 dx_ = self.time
247 dy_ = min(np.sqrt(self.time) * 0.6, 0.6)
248 if mode == 'compact':
249 x = host_x + hostdx_ - dx_
250 res = x, y, dx_, dy_
251 return res
253 def get_triangle(self, x: float, y: float, dx: float,
254 dy: float) -> List[List[float]]:
255 """Return the three triangle vertices for a send block pointing to the right."""
256 return [[x, y - dy / 2], [x, y + dy / 2], [x + dx, y]]
258 # pylint: disable=W0613
259 def draw_comm(self, ax: plt.Axes, *args: object, **kwargs: object) -> None:
260 """Draw the send→receive connector arrow between paired blocks on ``ax``."""
261 index_from, index_to = (kwargs.get('index_from', 0), kwargs.get('index_to', 0))
262 equal_wide, mode = (kwargs.get('equal_wide', False), kwargs.get('mode', 'compact'))
263 x, y, dx, _ = self.loc_size(index_from, equal_wide, mode)
264 x_, y_, dx_, _ = self.dual.loc_size(index_to, equal_wide, mode)
265 ax.annotate(None, xy=(x_ - dx_ / 2, y_), xytext=(x + dx / 2, y),
266 arrowprops={"ec": 'grey', "arrowstyle": '->', "shrinkA": 2, "shrinkB": 2})
268 @dfs_builder(True)
269 def build_with_comm(self) -> None:
270 r"""Build pipeline timeline with comm blocks and dependency."""
271 self.dual.depend_left.build_with_comm()
272 self.depend_left.build_with_comm()
273 self.start = max(self.depend_left.end, self.dual.depend_left.end)
274 self.end = self.start + self.time
276 def loop(self, comm: bool = False) -> List["BlockSim"]:
277 """Delegate to :meth:`comm_loop` when ``comm`` is ``True``, else fall back to base."""
278 if comm:
279 return self.comm_loop()
280 return super().loop(comm)
282 def comm_loop(self) -> list[BlockSim]:
283 r"""recursively check comm loop"""
284 if self.flag and not self.in_queue:
285 return []
286 res = []
287 if self.in_queue:
288 loop = [self]
289 block = self.father
290 while block.father and block is not self:
291 block = block.father
292 loop.append(block)
293 return loop
294 self.flag = True
295 self.in_queue = True
296 depends = [self.dual.depend_left, self.depend_left]
297 for dep in depends:
298 if dep:
299 dep.father = self
300 res.extend(dep.comm_loop())
301 dep.father = None
302 self.in_queue = False
303 return res
306@dataclass
307class RecBlockSim(CommBlockSim):
308 r"""sim receive comm block"""
309 type: str = 'r'
310 _color = '32'
312 def loc_size(self, x: float = 0, equal_wide: bool = False,
313 mode: str = 'compact') -> Tuple[float, float, float, float]:
314 """Return ``(x, y, dx, dy)`` for a receive block relative to its compute host."""
315 host_x, _, _, _ = self.host.loc_size(x, equal_wide)
316 x, y, _, _ = super().loc_size(x, equal_wide)
317 dx_ = self.time
318 dy_ = min(np.sqrt(self.time) * 0.6, 0.6)
319 if mode == 'compact':
320 x = host_x
321 res = x, y, -dx_, -dy_
322 return res
324 def get_triangle(self, x: float, y: float, dx: float,
325 dy: float) -> List[List[float]]:
326 """Return the three triangle vertices for a receive block pointing to the left."""
327 return [[x, y], [x - dx, y + dy / 2], [x - dx, y - dy / 2]]
329 @dfs_builder(True)
330 def build_with_comm(self) -> None:
331 r"""Build pipeline timeline with comm blocks and dependency."""
332 self.dual.build_with_comm()
333 self.depend_left.build_with_comm()
334 self.start = max(self.depend_left.end, self.dual.start)
335 self.end = self.start + self.time