Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / shard / ops / parallel_topk.py: 98%

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

16Distributed implementation for TopK operator. 

17""" 

18 

19from copy import deepcopy 

20from typing import Tuple 

21 

22from .parallel_ops import DistributedOp 

23 

24 

25def _normalize_topk_args(input_tensor, k, dim=-1, largest=True, sorted_output=True): 

26 return (input_tensor, k, dim, largest, sorted_output), {} 

27 

28 

29class TopKDistributedOp(DistributedOp): 

30 """Distributed implementation for TopK operator.""" 

31 

32 def preprocess(self, args: tuple, kwargs: dict) -> tuple: 

33 """ 

34 Preprocess arguments for TopK operator. 

35 

36 Args: 

37 args (tuple): Input arguments, first element is the input tensor. 

38 kwargs (dict): Keyword arguments. 

39 

40 Returns: 

41 tuple: (local_args, local_kwargs, cache_values) 

42 """ 

43 args, kwargs = _normalize_topk_args(*args, **kwargs) 

44 input_tensor = args[0] 

45 k = args[1] 

46 dim = args[2] 

47 largest = args[3] 

48 sorted_flag = args[4] 

49 

50 local_args = (input_tensor.to_local(), k, dim, largest, sorted_flag) 

51 local_kwargs = {} 

52 cache_values = [input_tensor.layout, dim] 

53 return local_args, local_kwargs, cache_values 

54 

55 def infer_layout(self, cache_values: list) -> Tuple[tuple, None]: 

56 """ 

57 Infer output layouts for TopK operator. 

58 

59 TopK: values, indices = topk(input, k, dim) 

60 

61 Rules: 

62 1. Input must not have Partial status. 

63 2. dim must be an integer within the valid range [-ndim, ndim-1]. 

64 3. The topk dimension must not be sharded (including StridedShard multi-axis mappings). 

65 4. Both values and indices output layouts are identical to the input layout. 

66 

67 Args: 

68 cache_values (list): [input_layout, dim] where dim is the topk dimension. 

69 

70 Returns: 

71 tuple: ((values_layout, indices_layout), None) 

72 

73 Raises: 

74 ValueError: If input has Partial status, dim is out of range, or the topk dimension 

75 is sharded. 

76 """ 

77 layout = cache_values[0] 

78 dim = cache_values[1] 

79 

80 if not self._allow_partial_inputs: 

81 self._check_partial_inputs([layout]) 

82 

83 if dim is None: 

84 dim = -1 

85 if not isinstance(dim, int): 

86 raise ValueError( 

87 f"For {self.op_name}, dimension should be int, but got {type(dim)}" 

88 ) 

89 

90 alias_map = layout.alias_tensor_map 

91 ndim = len(alias_map) 

92 

93 original_dim = dim 

94 if dim < 0: 

95 dim += ndim 

96 if not 0 <= dim < ndim: 

97 raise ValueError( 

98 f"For {self.op_name}, dimension out of range " 

99 f"(expected to be in range of [{-ndim}, {ndim - 1}], but got {original_dim})" 

100 ) 

101 

102 # Check if the topk dimension is sharded. 

103 # In alias_tensor_map, "None" means Replicate (not sharded); any other value implies sharding. 

104 mapping = alias_map[dim] 

105 if isinstance(mapping, (list, tuple)): 

106 is_sharded = any(m != "None" for m in mapping) 

107 else: 

108 is_sharded = mapping != "None" 

109 

110 if is_sharded: 

111 raise ValueError( 

112 f"For {self.op_name}, topk along a sharded dimension " 

113 f"(dim {dim} mapped to {mapping}) is not supported. " 

114 f"Please redistribute the tensor to Replicate on this dimension before topk." 

115 ) 

116 

117 return ((deepcopy(layout), deepcopy(layout)), None) 

118