Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / hyper_offload / planning / greedy.py: 99%

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"""Greedy planner for storage residency actions.""" 

16 

17from __future__ import annotations 

18 

19import logging 

20import pprint 

21from dataclasses import dataclass 

22 

23from hyper_parallel.auto_parallel.hyper_offload.ir.schedule import ResidencyActionType, ResidencySchedule 

24from hyper_parallel.auto_parallel.hyper_offload.ir.trace import ( 

25 AccessKind, 

26 ActivationTrace, 

27 StorageAccess, 

28) 

29 

30logger = logging.getLogger(__name__) 

31 

32 

33@dataclass(frozen=True) 

34class _EvictionCandidate: 

35 """An open interval between two adjacent accesses.""" 

36 

37 storage_id: int 

38 size_bytes: int 

39 copy_start: StorageAccess 

40 release_start: StorageAccess 

41 end: StorageAccess 

42 

43 @property 

44 def distance(self) -> int: 

45 """Return the distance.""" 

46 return self.end.op_id - self.release_start.op_id 

47 

48 @property 

49 def start_interior(self) -> int: 

50 """Return the start interior.""" 

51 return self.release_start.op_id + 1 

52 

53 @property 

54 def end_interior(self) -> int: 

55 """Return the end interior.""" 

56 return self.end.op_id - 1 

57 

58 

59@dataclass 

60class _Footprint: 

61 """Resident-memory footprint and derived eviction data for one storage.""" 

62 

63 size_bytes: int 

64 first_op_id: int 

65 last_op_id: int 

66 candidacies: list[_EvictionCandidate] 

67 last_access: StorageAccess 

68 

69 

70class GreedyResidencyPlanner: 

71 """Plan residency by evicting long access gaps until the budget is met.""" 

72 

73 def build(self, trace: ActivationTrace) -> ResidencySchedule: 

74 """Build the residency schedule from an activation trace. 

75 

76 Args: 

77 trace: Activation trace with ops, storage sizes and memory limit. 

78 

79 Returns: 

80 A :class:`ResidencySchedule` with pre/post eviction actions. 

81 """ 

82 schedule = ResidencySchedule() 

83 limit = trace.memory_limit_bytes if trace.memory_limit_bytes is not None else float("inf") 

84 if not trace.ops or limit == float("inf"): 

85 return schedule 

86 

87 accesses_by_storage = self._group_accesses_by_storage(trace) 

88 

89 resident_bytes, candidates, releasable = self._build_resident_footprint( 

90 trace, accesses_by_storage 

91 ) 

92 

93 selected = self._greedy_select(candidates, resident_bytes, limit) 

94 

95 for candidate in selected: 

96 self._add_eviction_actions(schedule, candidate) 

97 self._emit_release_actions(schedule, selected, trace, releasable) 

98 

99 peak = max(resident_bytes, default=0) 

100 self._log_planning_summary(trace, schedule, selected, candidates, peak, limit) 

101 

102 return schedule 

103 

104 # ------------------------------------------------------------------ 

105 # Phase 1: group accesses by storage 

106 # ------------------------------------------------------------------ 

107 

108 @staticmethod 

109 def _group_accesses_by_storage( 

110 trace: ActivationTrace, 

111 ) -> dict[int, list[StorageAccess]]: 

112 """Return a mapping from storage id to its ordered list of accesses.""" 

113 accesses_by_storage: dict[int, list[StorageAccess]] = {} 

114 for op in trace.ops: 

115 for access in op.accesses: 

116 accesses_by_storage.setdefault(access.storage_id, []).append(access) 

117 return accesses_by_storage 

118 

119 # ------------------------------------------------------------------ 

120 # Phase 2: compute resident footprint, eviction candidates & releasable 

121 # ------------------------------------------------------------------ 

122 

123 @classmethod 

124 def _build_resident_footprint( 

125 cls, 

126 trace: ActivationTrace, 

127 accesses_by_storage: dict[int, list[StorageAccess]], 

128 ) -> tuple[list[int], list[_EvictionCandidate], dict[int, StorageAccess]]: 

129 """Compute the per-op resident-memory array and enumerate eviction candidates. 

130 

131 Args: 

132 trace: Activation trace. 

133 accesses_by_storage: Per-storage access lists (from :meth:`_group_accesses_by_storage`). 

134 

135 Returns: 

136 A triple ``(resident_bytes, candidates, releasable)`` where 

137 *resident_bytes* is indexed by op id, 

138 *candidates* are sorted for greedy selection, and 

139 *releasable* maps each storage id to its last access. 

140 """ 

141 max_op_id = max( 

142 (access.op_id for op in trace.ops for access in op.accesses), 

143 default=0, 

144 ) 

145 resident_bytes = [0] * (max_op_id + 1) 

146 unsorted_candidates: list[_EvictionCandidate] = [] 

147 releasable: dict[int, StorageAccess] = {} 

148 

149 for sid, accesses in accesses_by_storage.items(): 

150 size_bytes = trace.storage_sizes.get(sid) 

151 if size_bytes is None or not accesses: 

152 continue 

153 

154 footprint = cls._process_storage_accesses(sid, size_bytes, accesses) 

155 for op_id in range(footprint.first_op_id, footprint.last_op_id + 1): 

156 resident_bytes[op_id] += size_bytes 

157 

158 releasable[sid] = footprint.last_access 

159 unsorted_candidates.extend(footprint.candidacies) 

160 

161 unsorted_candidates.sort( 

162 key=lambda candidate: ( 

163 candidate.distance, 

164 candidate.size_bytes, 

165 -candidate.release_start.op_id, 

166 ), 

167 reverse=True, 

168 ) 

169 

170 return resident_bytes, unsorted_candidates, releasable 

171 

172 @classmethod 

173 def _process_storage_accesses( 

174 cls, 

175 sid: int, 

176 size_bytes: int, 

177 accesses: list[StorageAccess], 

178 ) -> _Footprint: 

179 """Build the per-storage footprint: eviction candidates and extent info. 

180 

181 Args: 

182 sid: Storage id. 

183 size_bytes: Size of the storage in bytes. 

184 accesses: Ordered access list for this storage. 

185 

186 Returns: 

187 A :class:`_Footprint` with size, extent, eviction candidates and 

188 last access. 

189 """ 

190 ordered = sorted(accesses, key=lambda access: access.op_id) 

191 first = ordered[0].op_id 

192 last = ordered[-1].op_id 

193 candidacies: list[_EvictionCandidate] = [] 

194 

195 for release_index, (release_start, end) in enumerate( 

196 zip(ordered, ordered[1:], strict=False) 

197 ): 

198 copy_start = cls._earliest_safe_copy_start(ordered, release_index) 

199 candidate = _EvictionCandidate(sid, size_bytes, copy_start, release_start, end) 

200 if candidate.start_interior <= candidate.end_interior: 

201 candidacies.append(candidate) 

202 

203 return _Footprint( 

204 size_bytes=size_bytes, 

205 first_op_id=first, 

206 last_op_id=last, 

207 candidacies=candidacies, 

208 last_access=ordered[-1], 

209 ) 

210 

211 # ------------------------------------------------------------------ 

212 # Phase 3: greedy selection 

213 # ------------------------------------------------------------------ 

214 

215 @staticmethod 

216 def _greedy_select( 

217 candidates: list[_EvictionCandidate], 

218 resident_bytes: list[int], 

219 limit: float, 

220 ) -> list[_EvictionCandidate]: 

221 """Select eviction candidates greedily by longest gap first. 

222 

223 Args: 

224 candidates: Pre-sorted candidates (longest gap first). 

225 resident_bytes: Mutable per-op resident bytes (modified in-place). 

226 limit: Memory budget in bytes. 

227 

228 Returns: 

229 List of selected candidates. 

230 """ 

231 selected: list[_EvictionCandidate] = [] 

232 for candidate in candidates: 

233 if max(resident_bytes, default=0) <= limit: 

234 break 

235 

236 affected_ops = range(candidate.start_interior, candidate.end_interior + 1) 

237 if not any(resident_bytes[op_id] > limit for op_id in affected_ops): 

238 continue 

239 

240 selected.append(candidate) 

241 for op_id in affected_ops: 

242 resident_bytes[op_id] -= candidate.size_bytes 

243 

244 return selected 

245 

246 # ------------------------------------------------------------------ 

247 # Phase 4: emit release actions 

248 # ------------------------------------------------------------------ 

249 

250 @staticmethod 

251 def _emit_release_actions( 

252 schedule: ResidencySchedule, 

253 selected: list[_EvictionCandidate], 

254 trace: ActivationTrace, 

255 releasable: dict[int, StorageAccess], 

256 ) -> None: 

257 """Emit release-device and (conditional) release-host post-actions. 

258 

259 Args: 

260 schedule: Schedule to mutate. 

261 selected: Eviction candidates chosen by greedy selection. 

262 trace: Original activation trace (for ``retained_sids``). 

263 releasable: Mapping from storage id to its last access. 

264 """ 

265 evicted_sids = {candidate.storage_id for candidate in selected} 

266 for sid, access in releasable.items(): 

267 if sid in trace.retained_sids: 

268 continue 

269 schedule.add_post( 

270 access.op_id, 

271 ResidencyActionType.RELEASE_DEVICE, 

272 sid, 

273 ) 

274 if sid in evicted_sids: 

275 schedule.add_post( 

276 access.op_id, 

277 ResidencyActionType.RELEASE_HOST, 

278 sid, 

279 ) 

280 

281 # ------------------------------------------------------------------ 

282 # Logging helpers 

283 # ------------------------------------------------------------------ 

284 

285 @staticmethod 

286 def _log_planning_summary( 

287 trace: ActivationTrace, 

288 schedule: ResidencySchedule, 

289 selected: list[_EvictionCandidate], 

290 candidates: list[_EvictionCandidate], 

291 peak: int, 

292 limit: float, 

293 ) -> None: 

294 """Log the planning result at info and debug levels.""" 

295 logger.info( 

296 "Residency planner selected %d/%d gaps; simulated peak %.2f MiB, budget %.2f MiB", 

297 len(selected), 

298 len(candidates), 

299 peak / 1024**2, 

300 limit / 1024**2, 

301 ) 

302 logger.debug( 

303 "Schedule: %d pre-actions, %d post-actions", 

304 sum(len(v) for v in schedule.pre.values()), 

305 sum(len(v) for v in schedule.post.values()), 

306 ) 

307 GreedyResidencyPlanner._log_schedule_by_op(trace, schedule) 

308 

309 @staticmethod 

310 def _log_schedule_by_op( 

311 trace: ActivationTrace, 

312 schedule: ResidencySchedule, 

313 ) -> None: 

314 """Log the schedule details for each op (debug-level).""" 

315 for idx, op in enumerate(trace.ops): 

316 logger.debug(" OP %d name=%s", idx, op.name) 

317 for access in op.accesses: 

318 logger.debug( 

319 " access sid=%d %s", 

320 access.storage_id, 

321 access.kind.name, 

322 ) 

323 

324 pre_actions = schedule.pre_actions(idx) 

325 if pre_actions: 

326 logger.debug( 

327 " PRE %s", 

328 GreedyResidencyPlanner._format_actions(pre_actions, " PRE "), 

329 ) 

330 

331 post_actions = schedule.post_actions(idx) 

332 if post_actions: 

333 logger.debug( 

334 " POST %s", 

335 GreedyResidencyPlanner._format_actions(post_actions, " POST "), 

336 ) 

337 

338 @staticmethod 

339 def _earliest_safe_copy_start(accesses: list[StorageAccess], release_index: int) -> StorageAccess: 

340 """Return the earliest D2H point that still captures current data. 

341 

342 D2H may move before later reads because they do not mutate storage. 

343 It must not move before the latest write at or before the release 

344 point, otherwise the host copy could become stale. 

345 """ 

346 for index in range(release_index, -1, -1): 

347 if accesses[index].kind == AccessKind.WRITE: 

348 return accesses[index] 

349 return accesses[0] 

350 

351 @staticmethod 

352 def _format_actions(actions, prefix: str) -> str: 

353 items = [(action.kind.name, action.storage_id) for action in actions] 

354 formatted = pprint.pformat(items, compact=True, width=80) 

355 return formatted.replace("\n ", "\n" + " " * len(prefix)) 

356 

357 @staticmethod 

358 def _add_eviction_actions(schedule: ResidencySchedule, candidate: _EvictionCandidate) -> None: 

359 """Add pre/post schedule actions for a single eviction candidate.""" 

360 schedule.add_post( 

361 candidate.copy_start.op_id, 

362 ResidencyActionType.COPY_D2H, 

363 candidate.copy_start.storage_id, 

364 ) 

365 schedule.add_post( 

366 candidate.release_start.op_id, 

367 ResidencyActionType.RELEASE_DEVICE, 

368 candidate.release_start.storage_id, 

369 ) 

370 schedule.add_pre( 

371 candidate.end.op_id, 

372 ResidencyActionType.COPY_H2D, 

373 candidate.end.storage_id, 

374 )