Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / sapp_ppb / simulator / plot_manager.py: 80%
158 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"""Matplotlib canvas manager that renders the pipeline simulator timeline & memory plots."""
16from __future__ import annotations
18from collections.abc import Iterable
19from typing import List, Optional
21import matplotlib.pyplot as plt
22import numpy as np
23from matplotlib.transforms import ScaledTranslation
25from hyper_parallel.auto_parallel.sapp_ppb.simulator.sim_block import BlockSim, MicroBlockSim
28class PlotMgr:
29 """Holds a matplotlib figure and its axes, and provides the simulator draw helpers."""
31 # pylint: disable=W0613
32 def __init__(self, *args: object, num_plots: int = 2, ax_type: object = 'block',
33 subplot_args: Optional[List[int]] = None,
34 sub_fig: Optional[plt.Figure] = None, **kwargs: object) -> None:
35 """Create ``num_plots`` sub-axes on a new or reused matplotlib figure.
37 Args:
38 num_plots: Number of stacked sub-plots.
39 ax_type: Per-axis type label (or a single value replicated to each axis).
40 subplot_args: Optional explicit ``add_subplot`` specifiers; must have length
41 ``>= num_plots``.
42 sub_fig: Reuse this figure if given, otherwise create one of ``figsize``.
43 """
44 if sub_fig:
45 self.fig = sub_fig
46 else:
47 self.fig = plt.figure(figsize=kwargs.get('figsize', (12, 8)))
48 self.fig.subplots_adjust(wspace=0, hspace=0.4)
49 ax_type = ax_type if isinstance(ax_type, (list, tuple)) else [ax_type] * num_plots
50 self.ax: List[plt.Axes] = []
51 for i in range(num_plots):
52 if subplot_args is None:
53 self.ax.append(self.fig.add_subplot(num_plots * 100 + 10 + i + 1))
54 elif isinstance(subplot_args, Iterable) and len(subplot_args) >= num_plots:
55 self.ax.append(self.fig.add_subplot(subplot_args[i]))
56 else:
57 raise ValueError(f"Unsupported subplot_args format: {subplot_args}")
58 self.msg = ""
60 def _set_block_ax(self, ax: plt.Axes, pp: int) -> None:
61 """Configure one block-timeline axis (title, y ticks, y limits)."""
62 ax.set_title("Pipeline Flow Timeline")
63 ax.set_yticks(range(pp), [f"stage {p}" for p in range(pp)])
64 for tick in ax.get_yticklabels():
65 tick.set_verticalalignment('top')
66 tick.set_transform(
67 tick.get_transform() + ScaledTranslation(0, 0.05 - 1 / pp, self.fig.dpi_scale_trans))
68 tick.set_fontsize(12)
69 ax.set_ylim(0, pp)
70 ax.invert_yaxis()
72 def _get_block_indices(self, blocks: List[List[MicroBlockSim]],
73 mode: str = 'compact',
74 equal_wide: bool = False) -> List[np.ndarray]:
75 """Return per-stage cumulative x-coordinates suitable for drawing ``blocks``."""
76 if mode not in ['compact', 'joint', 'timeline']:
77 raise ValueError(f"Get unsupported draw mode: {mode}")
78 if mode == 'timeline' and not blocks[-1][-1].finish:
79 raise ValueError("Block building should be finished before drawing timeline")
80 block_index: List[np.ndarray] = []
81 for block_p in blocks:
82 inds: List[float] = []
83 for block in block_p:
84 if mode == 'compact':
85 if block.type == 'c':
86 inds.append(1 if equal_wide else block.time)
87 else:
88 inds.append(0)
89 elif mode == 'joint':
90 if block.type == 'c':
91 inds.append(1 if equal_wide else block.time)
92 else:
93 inds.append(block.time)
94 else:
95 inds.append(1)
96 inds.insert(0, 0)
97 inds = np.cumsum(inds)
98 block_index.append(inds)
99 return block_index
101 def draw_block(self, block_index: List[np.ndarray], blocks: List[List[MicroBlockSim]],
102 ax_index: int = 0, equal_wide: bool = False,
103 width: float = 1, phase: bool = False) -> "PlotMgr":
104 """Draw all compute blocks onto ``self.ax[ax_index]``."""
105 for p, block_p in enumerate(blocks):
106 for b, block in enumerate(block_p):
107 if block.type == 'c':
108 block.draw(self.ax[ax_index], index=block_index[p][b],
109 equal_wide=equal_wide, width=width, phase=phase)
110 return self
112 def draw_comm(self, block_index: List[np.ndarray], blocks: List[List[MicroBlockSim]],
113 ax_index: int = 0, equal_wide: bool = False,
114 mode: str = 'compact') -> "PlotMgr":
115 """Draw send/receive comm blocks onto ``self.ax[ax_index]``."""
116 for p, block_p in enumerate(blocks):
117 for b, block in enumerate(block_p):
118 if block.type == 'c' and mode == 'compact':
119 if block.send_block:
120 block.send_block.draw(self.ax[ax_index], index=block_index[p][b],
121 equal_wide=equal_wide)
122 if block.rec_block:
123 block.rec_block.draw(self.ax[ax_index], index=block_index[p][b],
124 equal_wide=equal_wide)
125 elif block.type in ['s', 'r'] and mode in ['joint', 'timeline']:
126 block.draw(self.ax[ax_index], index=block_index[p][b],
127 equal_wide=equal_wide, mode=mode)
128 return self
130 def draw_connect(self, block_index: List[np.ndarray], blocks: List[List[MicroBlockSim]],
131 ax_index: int = 0, equal_wide: bool = False,
132 mode: str = 'compact') -> "PlotMgr":
133 """Draw the arrows that connect each send block to its matching receive block."""
134 for p, block_p in enumerate(blocks):
135 for b, block in enumerate(block_p):
136 if block.type == 'c' and mode == 'compact' and block.send_block:
137 dual_p = block.send_block.dual.stage
138 dual_ind = blocks[dual_p].index(block.send_block.dual.host)
139 block.send_block.draw_comm(
140 self.ax[ax_index], index_from=block_index[p][b],
141 index_to=block_index[dual_p][dual_ind],
142 equal_wide=equal_wide, mode=mode)
143 elif block.type == 's' and mode in ['joint', 'timeline']:
144 dual_p = block.dual.stage
145 dual_ind = blocks[dual_p].index(block.dual)
146 block.draw_comm(
147 self.ax[ax_index], index_from=block_index[p][b],
148 index_to=block_index[dual_p][dual_ind],
149 equal_wide=equal_wide, mode=mode)
150 return self
152 def draw(self, blocks: List[List[MicroBlockSim]], ax_index: int = 0,
153 comm: bool = False, connect: bool = False,
154 equal_wide: bool = False, mode: str = 'compact',
155 phase: bool = False) -> "PlotMgr":
156 """Draw the full pipeline timeline: blocks, comm layer and connect arrows."""
157 pp = len(blocks)
158 block_index = self._get_block_indices(blocks, mode=mode, equal_wide=equal_wide)
159 width = max(np.max(block_index[p]) for p in range(pp)) if blocks[0][-1].end is None \
160 else max(blocks[p][-1].end for p in range(pp))
161 plot_mgr = self.draw_block(block_index, blocks, ax_index, equal_wide, width, phase=phase)
162 if plot_mgr is not self:
163 raise RuntimeError("Unexpected draw result.")
164 if comm:
165 plot_mgr = self.draw_comm(block_index, blocks, ax_index, equal_wide, mode)
166 if plot_mgr is not self:
167 raise RuntimeError("Unexpected draw result.")
168 if connect:
169 plot_mgr = self.draw_connect(block_index, blocks, ax_index, equal_wide, mode)
170 if plot_mgr is not self:
171 raise RuntimeError("Unexpected draw result.")
172 self._set_block_ax(self.ax[ax_index], pp)
173 self.ax[ax_index].set_xlim(0, width)
174 self.ax[ax_index].set_xticks(np.linspace(0, width, 8))
175 return self
177 def draw_loop(self, blocks: List[List[MicroBlockSim]], loop: List[BlockSim],
178 ax_index: int = 0, comm: bool = False, connect: bool = False,
179 equal_wide: bool = False) -> "PlotMgr":
180 """Highlight a dependency loop (non-comm) with red arrows and a textual trace."""
181 plot_mgr = self.draw(blocks, ax_index, comm, connect, equal_wide, phase=True)
182 if plot_mgr is not self:
183 raise RuntimeError("Unexpected draw result.")
184 block_index = self._get_block_indices(blocks, equal_wide=equal_wide)
185 msg = 'dependency loop: '
186 for b in range(len(loop) - 1):
187 p = loop[b].stage
188 ind = blocks[p].index(loop[b])
189 x1, y1, dx1, _ = loop[b].loc_size(block_index[p][ind], equal_wide)
190 p = loop[b + 1].stage
191 ind = blocks[p].index(loop[b + 1])
192 x2, y2, dx2, _ = loop[b + 1].loc_size(block_index[p][ind], equal_wide)
193 msg = f'{msg} {loop[b].color_label} -> '
194 self.ax[ax_index].annotate(
195 None, xy=(x1 + dx1 / 2, y1), xytext=(x2 + dx2 / 2, y2),
196 arrowprops={"fc": 'white', "ec": 'r', "arrowstyle": 'simple',
197 "shrinkA": 5, "shrinkB": 5,
198 "connectionstyle": "arc3,rad=-0.1"})
199 self.msg = f'{msg} {loop[len(loop) - 1].color_label}'
200 return self
202 def draw_comm_loop(self, lines: List[List[BlockSim]], loop: List[BlockSim],
203 ax_index: int = 0) -> "PlotMgr":
204 """Highlight a dependency loop in the send-receive graph."""
205 draw_result = self.draw(lines, ax_index, True, True, True, 'joint', phase=True)
206 if draw_result is not self:
207 raise RuntimeError("Unexpected draw result.")
208 block_index = self._get_block_indices(lines, mode='joint', equal_wide=True)
209 msg = 'dependency loop: '
210 for b in range(len(loop) - 1):
211 p = loop[b].stage
212 ind = lines[p].index(loop[b])
213 x1, y1, dx1, _ = loop[b].loc_size(block_index[p][ind], True, 'joint')
214 p = loop[b + 1].stage
215 ind = lines[p].index(loop[b + 1])
216 x2, y2, dx2, _ = loop[b + 1].loc_size(block_index[p][ind], True, 'joint')
217 msg = f'{msg} {loop[b].color_label} -> '
218 self.ax[ax_index].annotate(
219 None, xy=(x1 + abs(dx1) / 2, y1), xytext=(x2 + abs(dx2) / 2, y2),
220 size=10,
221 arrowprops={"fc": 'white', "ec": 'r', "arrowstyle": 'simple',
222 "shrinkA": 3, "shrinkB": 3,
223 "connectionstyle": "arc3,rad=-0.1", "lw": 0.8})
224 self.msg = f'{msg} {loop[len(loop) - 1].color_label}'
225 return self
227 def draw_mem(self, block_mem_list: List[np.ndarray], ax_index: int = 0) -> "PlotMgr":
228 """Plot per-stage block memory curves on the axis at ``ax_index``."""
229 for p, block_mem in enumerate(block_mem_list):
230 self.ax[ax_index].plot((block_mem.T)[0], (block_mem.T)[1], label=f"stage-{p}")
231 self.ax[ax_index].set_title("Block Memory Timeline")
232 width = max(np.max((block_mem.T)[0]) for block_mem in block_mem_list)
233 height = max(np.max((block_mem.T)[1]) for block_mem in block_mem_list)
234 self.ax[ax_index].set_xlim(
235 0, max(np.max((block_mem.T)[0]) for block_mem in block_mem_list))
236 self.ax[ax_index].set_xticks(np.linspace(0, width, 8))
237 self.ax[ax_index].set_yticks(np.linspace(0, height, 4))
238 return self
240 def draw_info(self, bubble_info: dict, mem_info: List[float]) -> None:
241 """Draw the bubble / peak-memory annotation lines at the top & bottom of the figure."""
242 info_list = [f'{k} bubble: {v:.4f}' for k, v in bubble_info.items()]
243 self.fig.text(0.5, 0.5, ', '.join(info_list), ha='center', va='center',
244 fontdict={'fontsize': 13, 'weight': 'medium'}, color='C3')
245 info_list = [f"{v:.0f}" for v in mem_info]
246 self.fig.text(0.5, 0.05, f"peak memory: {', '.join(info_list)}", ha='center', va='center',
247 fontdict={'fontsize': 10, 'weight': 'medium'}, color='C0')
249 def save(self, file_name: str) -> None:
250 """Save the figure to ``file_name``."""
251 self.fig.legend(bbox_to_anchor=(0.22, 0.45))
252 plt.savefig(file_name)
254 def show(self, file_name: Optional[str] = None) -> None:
255 """Display the figure interactively, optionally also saving it to ``file_name``."""
256 self.fig.legend(bbox_to_anchor=(0.22, 0.45))
257 if file_name is not None:
258 plt.savefig(file_name)
259 plt.show()