Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / sapp_nd / perf_estimation / getters.py: 100%

1# Copyright 2025 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"""Getters""" 

16 

17from copy import deepcopy 

18from hyper_parallel.auto_parallel.sapp_nd.nd.common.layer_type import LayerType 

19from hyper_parallel.auto_parallel.sapp_nd.nd.logger import perf_logger as logger 

20 

21 

22# Configs are updated depending on the type of the transformer layer 

23def get_layer_custom_configs(cfg): 

24 """Stores each configuration along with how many layers are affected by it 

25 in ascending order of execution in a forward pass 

26 """ 

27 

28 if cfg.layer_custom_config is None or any( 

29 func is None for (_, func) in cfg.layer_custom_config 

30 ): 

31 return [(cfg, cfg.n_lay)] 

32 

33 lccfgs = [] 

34 for nb_layers, func in cfg.layer_custom_config: 

35 lccfg = deepcopy(cfg) 

36 func(lccfg) 

37 lccfgs.append((lccfg, nb_layers)) 

38 

39 return lccfgs 

40 

41 

42def get_recomp_factor(lccfg, layer, op_name): 

43 """recomputation factor""" 

44 if layer == LayerType.FULL_REC_LAYER: 

45 return 1 

46 if layer == LayerType.NOT_REC_LAYER: 

47 return 0 

48 if layer == LayerType.SEL_REC_LAYER: 

49 return getattr(lccfg.rec_op, op_name, 0) 

50 logger.warning("Unrecognized recompute type %s", layer) 

51 return 0 

52 

53 

54def get_table_quantity(lccfg, table, layer, with_recomp): 

55 """op compute load from given table""" 

56 qt_layer = 0 

57 for op, quantity in table.items(): 

58 op_name = op[2:] 

59 

60 qt_layer += ( 

61 (1 + with_recomp * get_recomp_factor(lccfg, layer, op_name)) 

62 * getattr(lccfg, op) 

63 * quantity 

64 ) 

65 return qt_layer