Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_multicore / modules / mega_moe / _gmm_tiling.py: 95%

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 tiling template expansion and dimension patching.""" 

16 

17import ctypes 

18 

19from hyper_multicore._core_spec import DEFAULT_CORE_SPEC 

20 

21 

22def _expand_gmm_string(csv: str, repeat: int = DEFAULT_CORE_SPEC.num_cube_cores) -> bytes: 

23 """Parse comma-separated uint32 string and tile it `repeat` times.""" 

24 vals = [int(x) for x in csv.split(',')] 

25 arr = (ctypes.c_uint32 * (len(vals) * repeat))() 

26 for i in range(len(vals) * repeat): 

27 arr[i] = vals[i % len(vals)] 

28 return bytes(arr) 

29 

30 

31def _find_protected_positions(vals, active_map, split_value, is_weight_grad): 

32 """Return index positions that must not be replaced by dim_map substitution. 

33 

34 When split_value collides with a K/N dim in active_map, the M-tile or 

35 K-tile positions (depending on GMM type) must be shielded. 

36 T is located via the 393216 anchor which sits at T+19. 

37 """ 

38 if split_value is None or split_value not in active_map: 

39 return set() 

40 for i in range(350, len(vals)): 

41 if vals[i] == 393216: # anchor always at t_base+19 

42 t_base = i - 19 

43 if is_weight_grad: 

44 return {t_base + 3, t_base + 4, t_base + 7} 

45 return {14, t_base + 1, t_base + 5} 

46 return set() 

47 

48 

49def _patch_gmm_dims(csv: str, dim_map: dict, num_groups: int, 

50 default_groups: int = 8, 

51 split_value: int = None, 

52 num_cube_cores: int = DEFAULT_CORE_SPEC.num_cube_cores, 

53 default_num_cores: int = 24) -> str: 

54 """Replace matrix dimension values in a GMM tiling CSV string. 

55 

56 dim_map: {old_int_value: new_int_value} for K/N dimension replacements. 

57 Entries where old == new are no-ops and can be omitted. 

58 num_groups: replacement for default_groups; only applied at index > 350. 

59 split_value: M tile size (table key). When it equals a K/N value in dim_map, 

60 we must protect M-tile positions from accidental replacement. 

61 Two GMM types detected via vals[10]: 

62 0 → activation GMM: M tile at {14, T+1, T+5} 

63 2 → weight-grad GMM: split (K-tile) at {T+3, T+4, T+7} 

64 T is found via the 393216 anchor fixed at T+19. 

65 num_cube_cores: number of Cube cores used by the operator. 

66 Replaces the hardcoded usedCoreNum/blockDim fields in the CSV: 

67 vals[1] — header blockDim field 

68 i > 350 — tail usedCoreNum and any derived parallelism fields 

69 """ 

70 active_map = {k: v for k, v in dim_map.items() if k != v} 

71 groups_changed = num_groups != default_groups 

72 cores_changed = num_cube_cores != default_num_cores 

73 if not active_map and not groups_changed and not cores_changed: 

74 return csv # fast path: nothing to do 

75 

76 vals = list(map(int, csv.split(','))) 

77 is_weight_grad = vals[10] == 2 

78 

79 # Patch groupNum and blockDim before dimension substitution. 

80 for index, old_value, new_value in ((0, default_groups, num_groups), 

81 (1, default_num_cores, num_cube_cores)): 

82 if vals[index] == old_value: 

83 vals[index] = new_value 

84 

85 protected = _find_protected_positions(vals, active_map, split_value, is_weight_grad) 

86 

87 for i, v in enumerate(vals): 

88 if i in protected: 

89 continue 

90 if v in active_map: 

91 vals[i] = active_map[v] 

92 elif cores_changed and i > 350 and v == default_num_cores: 

93 vals[i] = num_cube_cores # covers T+0 (usedCoreNum) and derived fields 

94 return ','.join(map(str, vals))