Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_multicore / tasks / utils.py: 73%

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""" 

16Shared utilities for all FillConfig implementations. 

17 

18Counter helpers (advance_tsv_*) mirror the C++ tsv state advancement. 

19add_terminate / add_dynamic_data are called from gen_runtime_data after 

20the main fill loop. 

21""" 

22from hyper_multicore.scheduler.config import ( 

23 TensorDescC, TaskDescC, RuntimeConfigC, 

24 TaskType, DynamicType, 

25 MAX_TENSOR_DIMS, 

26) 

27 

28 

29# ── TensorDescC quick constructor ───────────────────────────────────────────── 

30 

31def make_tensor_desc( 

32 tensor_type: int = 0, 

33 data_type: int = 2, 

34 input_position: int = 0, 

35 base_ptr_offset: int = 0, 

36 transpose_flag: int = 0, 

37 dynamic_shape: int = 0, 

38 dynamic_dim: int = 0, 

39 shape: list = None, 

40) -> TensorDescC: 

41 """Build a TensorDescC with the given fields.""" 

42 t = TensorDescC() 

43 t.tensor_type = tensor_type 

44 t.data_type = data_type 

45 t.input_position = input_position 

46 t.base_ptr_offset = base_ptr_offset 

47 t.transpose_flag = transpose_flag 

48 t.dynamic_shape = dynamic_shape 

49 t.dynamic_dim = dynamic_dim 

50 if shape: 

51 for k, v in enumerate(shape[:MAX_TENSOR_DIMS]): 

52 t.dim[k] = v 

53 return t 

54 

55 

56# ── tsv counter helpers ─────────────────────────────────────────────────────── 

57 

58def advance_tsv_vector(tsv, task_num: int, event_group_size: int) -> None: 

59 """Advance counters after a vector op with event grouping.""" 

60 tsv.pre_task_num += task_num 

61 tsv.pre_vector_task_num += task_num 

62 tsv.pre_pre_event_num = tsv.pre_event_num 

63 tsv.pre_event_num += event_group_size 

64 

65 

66def advance_tsv_cube(tsv, task_num: int, event_group_size: int) -> None: 

67 """Advance counters after a cube op with event grouping.""" 

68 tsv.pre_task_num += task_num 

69 tsv.pre_cube_task_num += task_num 

70 tsv.pre_pre_event_num = tsv.pre_event_num 

71 tsv.pre_event_num += event_group_size 

72 

73 

74def advance_tsv_cube_only(tsv, task_num: int) -> None: 

75 """Advance only task/cube counters (no event update). Used by bwd GMM1.""" 

76 tsv.pre_task_num += task_num 

77 tsv.pre_cube_task_num += task_num 

78 

79 

80def advance_tsv_vector_only(tsv, task_num: int) -> None: 

81 """Advance only task/vector counters (no event update). Used by bwd combine.""" 

82 tsv.pre_task_num += task_num 

83 tsv.pre_vector_task_num += task_num 

84 

85 

86# ── Post-fill utilities ─────────────────────────────────────────────────────── 

87 

88def add_terminate(cfg: RuntimeConfigC, tsv, trigger_count: int) -> None: 

89 """ 

90 Append a terminate task to cfg. 

91 

92 trigger_count: 

93 forward: combine_op.task_num // tsv.ep * tsv.ep 

94 backward: w2_grad_op.task_num + w1_grad_op.task_num + combine_op.task_num // tsv.ep * tsv.ep 

95 """ 

96 cfg.all_event_num_triggers[tsv.all_event_num] = trigger_count 

97 task = TaskDescC() 

98 task.task_type = TaskType.TASK_TERMINATE 

99 task.dependent_event = tsv.pre_pre_event_num + 1 

100 task.trigger_event = tsv.pre_event_num + 1 

101 cfg.all_event_num_triggers[task.trigger_event] = 1 

102 

103 cfg.all_tasks[tsv.pre_task_num] = task 

104 cfg.vector_task_indices[tsv.pre_vector_task_num] = tsv.pre_task_num 

105 cfg.task_index_num[1] += 1 

106 

107 

108def add_dynamic_data(cfg: RuntimeConfigC, tsv, dynamic_input_position: int) -> None: 

109 """ 

110 Write the dynamic data record into cfg. 

111 

112 dynamic_input_position: forward = 6, backward = 19. 

113 """ 

114 cfg.dynamic_data.dynamic_type = DynamicType.DYNAMIC_DSV3_MOE 

115 cfg.dynamic_data.dynamic_input_position = dynamic_input_position 

116 cfg.dynamic_data.dynamic_group_size = tsv.single_rank_expert_num 

117 cfg.dynamic_data.dynamic_max_seq_len = -1