Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / tensor_parallel / _ce_op_registry.py: 86%

1# Copyright 2025-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"""CE operator name registry. 

16 

17Set of operator names on the CE decomposition path; elements are platform.get_op_name results. 

18Allows extension via register_loss_parallel_op_names (for test injection). 

19""" 

20 

21from typing import FrozenSet, Set 

22 

23__all__ = [ 

24 "get_loss_parallel_op_names", 

25 "get_decomposed_ce_op_names", 

26 "register_loss_parallel_op_names", 

27 "is_loss_parallel_op", 

28 "is_decomposed_ce_op", 

29] 

30 

31 

32_BASE_CE_OP_NAMES: FrozenSet[str] = frozenset({ 

33 "cross_entropy", 

34 "torch_cross_entropy", 

35 "torch_nn_CrossEntropyLoss", 

36 "mint_nn_functional_cross_entropy", 

37 "CrossEntropy", 

38}) 

39 

40_DECOMPOSED_CE_OP_NAMES: FrozenSet[str] = frozenset({ 

41 "nll_loss", 

42 "log_softmax", 

43 "softmax", 

44 "torch_nll_loss", 

45 "torch_log_softmax", 

46 "torch_softmax", 

47 "nll_loss_forward", 

48 "nll_loss_backward", 

49 "log_softmax_backward", 

50 "softmax_backward", 

51 "mint_nn_functional_nll_loss", 

52 "mint_nn_functional_log_softmax", 

53 "mint_nn_functional_softmax", 

54 "NLLLoss", 

55 "LogSoftmax", 

56 "Softmax", 

57}) 

58 

59_EXTENDED_CE_OP_NAMES: Set[str] = set() 

60 

61 

62def get_loss_parallel_op_names() -> FrozenSet[str]: 

63 """Return CE-related operator name set. 

64 

65 Returns: 

66 FrozenSet[str]: Full set of operator names on CE decomposition path. 

67 """ 

68 return _BASE_CE_OP_NAMES | frozenset(_EXTENDED_CE_OP_NAMES) 

69 

70 

71def get_decomposed_ce_op_names() -> FrozenSet[str]: 

72 """Return decomposed CE-related operator name set. 

73 

74 Returns: 

75 FrozenSet[str]: Set of operator names that are part of CE decomposition 

76 but should not be called directly in loss_parallel context. 

77 """ 

78 return _DECOMPOSED_CE_OP_NAMES 

79 

80 

81# Backward compatibility aliases (will be deprecated in future) 

82LOSS_PARALLEL_OP_NAMES = get_loss_parallel_op_names 

83DECOMPOSED_CE_OP_NAMES = get_decomposed_ce_op_names 

84 

85 

86def register_loss_parallel_op_names(*names: str) -> None: 

87 """Extend CE operator name set (for test injection). 

88 

89 Args: 

90 *names: Operator names to register (platform.get_op_name results). 

91 """ 

92 _EXTENDED_CE_OP_NAMES.update(names) 

93 

94 

95def unregister_loss_parallel_op_names(*names: str) -> None: 

96 """Remove registered CE operator names (for test cleanup). 

97 

98 Args: 

99 *names: Operator names to remove. 

100 """ 

101 _EXTENDED_CE_OP_NAMES.difference_update(names) 

102 

103 

104def is_loss_parallel_op(op_name: str) -> bool: 

105 """Check if operator is a CE entry point (e.g., cross_entropy). 

106 

107 Args: 

108 op_name: Operator name (platform.get_op_name result). 

109 

110 Returns: 

111 bool: Whether operator is a CE entry point. 

112 """ 

113 return op_name in get_loss_parallel_op_names() 

114 

115 

116def is_decomposed_ce_op(op_name: str) -> bool: 

117 """Check if operator is a decomposed CE op (e.g., log_softmax, nll_loss). 

118 

119 Args: 

120 op_name: Operator name (platform.get_op_name result). 

121 

122 Returns: 

123 bool: Whether operator is a decomposed CE op. 

124 """ 

125 return op_name in get_decomposed_ce_op_names() 

126 

127 

128def clear_extended_op_names() -> None: 

129 """Clear extended operator name set (for test cleanup).""" 

130 _EXTENDED_CE_OP_NAMES.clear()