Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / hyper_offload / ir / schedule.py: 100%

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"""Scheduler output model for residency actions.""" 

16 

17from __future__ import annotations 

18 

19from dataclasses import dataclass, field 

20from enum import Enum, auto 

21 

22 

23class ResidencyActionType(Enum): 

24 """Runtime action emitted by the planner.""" 

25 

26 COPY_D2H = auto() 

27 COPY_H2D = auto() 

28 RELEASE_DEVICE = auto() 

29 RELEASE_HOST = auto() 

30 

31 

32@dataclass(frozen=True) 

33class ResidencyAction: 

34 """A scheduled runtime residency action.""" 

35 

36 op_id: int 

37 storage_id: int 

38 kind: ResidencyActionType 

39 

40 

41@dataclass 

42class ResidencySchedule: 

43 """Scheduler output indexed by op id.""" 

44 

45 pre: dict[int, list[ResidencyAction]] = field(default_factory=dict) 

46 post: dict[int, list[ResidencyAction]] = field(default_factory=dict) 

47 

48 def add_pre(self, op_id: int, kind: ResidencyActionType, storage_id: int) -> None: 

49 """Add pre.""" 

50 self.pre.setdefault(op_id, []).append(ResidencyAction(op_id, storage_id, kind)) 

51 

52 def add_post(self, op_id: int, kind: ResidencyActionType, storage_id: int) -> None: 

53 """Add post.""" 

54 self.post.setdefault(op_id, []).append(ResidencyAction(op_id, storage_id, kind)) 

55 

56 def pre_actions(self, op_id: int) -> list[ResidencyAction]: 

57 """Pre actions.""" 

58 return self.pre.get(op_id, []) 

59 

60 def post_actions(self, op_id: int) -> list[ResidencyAction]: 

61 """Post actions.""" 

62 return self.post.get(op_id, [])