Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_multicore / tasks / gmm.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"""GMM fill config: all GroupedMatMul variants (fwd up_proj/down_proj, bwd act/gate/w1/w2 grad).""" 

16from dataclasses import dataclass, field 

17from typing import Set 

18 

19from hyper_multicore._core_spec import DEFAULT_CORE_SPEC 

20from hyper_multicore.scheduler.config import ( 

21 TaskDescC, TensorDescC, RuntimeConfigC, 

22 TaskAiCoreType, TaskType, TaskSplitValue, 

23 MAX_TENSOR_DIMS, 

24) 

25from hyper_multicore.tasks.task_base import FillConfig 

26from hyper_multicore.scheduler.graph import OperatorNode 

27from hyper_multicore.tasks.utils import ( 

28 advance_tsv_cube, advance_tsv_cube_only, 

29) 

30 

31 

32@dataclass 

33class GmmFillConfig(FillConfig): 

34 """ 

35 GMM fill config covering all variants (fwd up_proj/down_proj, bwd act_grad/w2_grad/ 

36 gate_grad/w1_grad). Tensor attributes (tensor_type, dtype_size, is_dynamic, 

37 transpose, shape) are read from TensorSpec. 

38 

39 offset_inputs : Set[int] 

40 Input indices that receive base_ptr_offset and dynamic_shape. 

41 activation GMMs (up_proj, down_proj, act_grad, gate_grad): {0} 

42 weight-grad GMMs (w2_grad, w1_grad): {0, 1} 

43 

44 transpose_inputs : Set[int] 

45 Additional transposed operands, without changing shared TensorSpec objects. 

46 

47 rank_in_event : bool 

48 True → dependent_event adds single_rank_expert_num * rank_id. 

49 GMM1 (fwd/bwd), GMM4/w2_grad (bwd). 

50 False → no rank offset. 

51 GMM2 (fwd/bwd), GMM3/w1_grad (bwd). 

52 

53 global_trigger : bool 

54 False → trigger_event = pre_event_num + data_index + 1 (per-expert). 

55 GMM1/GMM2 (activation-gradient path). 

56 True → trigger_event = all_event_num (global). 

57 GMM4/GMM3 (weight-gradient path). 

58 

59 out_offset : bool 

60 True → out.base_ptr_offset = data_index * 4096 * shape[1]. 

61 False → 0 (weight-grad output writes to a dedicated buffer). 

62 

63 advance : "cube" | "cube_only" | "cube_custom" 

64 "cube" — standard cube advance (also advances events). 

65 "cube_only" — advances task/cube counters only (no event update). 

66 Used by: GMM1 bwd. 

67 "cube_custom" — manual: pre_task_num/pre_cube_task_num += task_num, 

68 pre_pre_event_num = pre_event_num, 

69 pre_event_num += event_delta. 

70 Used by: GMM4 bwd, GMM3 bwd. 

71 

72 event_delta : int 

73 Only used when advance="cube_custom". 

74 GMM4 bwd: single_rank_expert_num. GMM3 bwd: 1. 

75 

76 num_cube_cores : int 

77 Number of Cube cores. 

78 """ 

79 offset_inputs: Set[int] = field(default_factory=lambda: {0}) 

80 # Shared activation specs retain their producer's interpretation. 

81 transpose_inputs: Set[int] = field(default_factory=set) 

82 rank_in_event: bool = False 

83 global_trigger: bool = False 

84 out_offset: bool = True 

85 advance: str = "cube" # "cube" | "cube_only" | "cube_custom" 

86 event_delta: int = 0 # only used when advance="cube_custom" 

87 num_cube_cores: int = DEFAULT_CORE_SPEC.num_cube_cores 

88 

89 def fill(self, cfg: RuntimeConfigC, op: OperatorNode, tsv: TaskSplitValue) -> None: 

90 """Serialize operand interpretation and event dependencies. 

91 

92 Args: 

93 cfg: Runtime descriptor to populate. 

94 op: Operator supplying operands and split metadata. 

95 tsv: Current task indices and expert dimensions. 

96 """ 

97 task_num = op.task_num 

98 param = op.param_positions 

99 glist_j = len(op.inputs) - 1 # last input is always group_list 

100 

101 for i in range(task_num): 

102 data_index = i // self.num_cube_cores 

103 task = TaskDescC() 

104 task.task_type = TaskType.TASK_GROUPED_MATMUL 

105 task.task_aicore_type = TaskAiCoreType.TASK_AICORE_CUBE 

106 task.num_inputs = len(op.inputs) 

107 task.num_outputs = len(op.outputs) 

108 

109 for j, spec in enumerate(op.inputs): 

110 td = TensorDescC() 

111 td.input_position = param[j] 

112 td.dynamic_shape = 0 

113 td.base_ptr_offset = 0 

114 

115 if j == glist_j: # group_list: fixed type=0, dtype int64 

116 td.tensor_type = 0 

117 td.data_type = spec.dtype_size 

118 else: # x or weight: tensor_type from TensorSpec 

119 td.tensor_type = spec.tensor_type 

120 td.data_type = spec.dtype_size 

121 

122 if j in self.offset_inputs: 

123 td.base_ptr_offset = data_index * 4096 * spec.shape[1] 

124 td.dynamic_shape = int(spec.is_dynamic) 

125 

126 td.transpose_flag = int(spec.transpose or j in self.transpose_inputs) 

127 for k in range(min(len(spec.shape), MAX_TENSOR_DIMS)): 

128 td.dim[k] = spec.shape[k] 

129 task.inputs[j] = td 

130 

131 out_spec = op.outputs[0] 

132 out = TensorDescC() 

133 out.tensor_type = out_spec.tensor_type 

134 out.data_type = out_spec.dtype_size 

135 out.input_position = param[task.num_inputs] 

136 out.dynamic_shape = int(out_spec.is_dynamic) 

137 out.base_ptr_offset = (data_index * 4096 * out_spec.shape[1] 

138 if self.out_offset else 0) 

139 if self.out_offset: # C++ fills dims only for activation outputs 

140 for k in range(min(len(out_spec.shape), MAX_TENSOR_DIMS)): 

141 out.dim[k] = out_spec.shape[k] 

142 task.outputs[0] = out 

143 

144 dep_extra = (tsv.single_rank_expert_num * tsv.rank_id 

145 if self.rank_in_event else 0) 

146 task.dependent_event = tsv.pre_pre_event_num + dep_extra + data_index + 1 

147 if self.global_trigger: 

148 task.trigger_event = tsv.all_event_num 

149 else: 

150 task.trigger_event = tsv.pre_event_num + data_index + 1 

151 cfg.all_event_num_triggers[task.trigger_event] = self.num_cube_cores 

152 

153 task.task_index = i 

154 task.task_split_num = task_num 

155 task.task_split_value = op.split_value 

156 task.tiling_data_position = op.tiling_position 

157 

158 cfg.all_tasks[tsv.pre_task_num + i] = task 

159 cfg.cube_task_indices[tsv.pre_cube_task_num + i] = tsv.pre_task_num + i 

160 

161 cfg.task_index_num[0] += task_num 

162 if self.advance == "cube": 

163 advance_tsv_cube(tsv, task_num, event_group_size=task_num // self.num_cube_cores) 

164 elif self.advance == "cube_only": 

165 advance_tsv_cube_only(tsv, task_num) 

166 else: # "cube_custom" 

167 tsv.pre_task_num += task_num 

168 tsv.pre_cube_task_num += task_num 

169 tsv.pre_pre_event_num = tsv.pre_event_num 

170 tsv.pre_event_num += self.event_delta