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

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

16Distributed implementation for ChunkView operator. 

17""" 

18 

19from typing import Tuple 

20 

21from .parallel_ops import DistributedOp 

22 

23 

24def _normalize_chunk_view_args(input_tensor, chunks, dim=0): 

25 return (input_tensor, chunks, dim), {} 

26 

27 

28class ChunkViewDistributedOp(DistributedOp): 

29 """Distributed implementation for ChunkView operator.""" 

30 

31 @staticmethod 

32 def _calculate_output_count(dim_size, chunks): 

33 """Calculate the number of output chunks based on dimension size.""" 

34 if dim_size == 0: 

35 return chunks 

36 split_size = (dim_size + chunks - 1) // chunks 

37 output_num = max((dim_size + split_size - 1) // split_size, 1) 

38 return min(output_num, chunks) 

39 

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

41 """ 

42 Preprocess arguments for ChunkView operator. 

43 

44 Args: 

45 args (tuple): Input arguments containing the input tensor, chunks, and dim. 

46 kwargs (dict): Keyword arguments (none expected). 

47 

48 Returns: 

49 tuple: (local_args, local_kwargs, cache_values) 

50 """ 

51 args, kwargs = _normalize_chunk_view_args(*args, **kwargs) 

52 input_tensor, chunks, dim = args 

53 input_shape = input_tensor.shape 

54 

55 local_args = (input_tensor.to_local(), chunks, dim) 

56 local_kwargs = {} 

57 

58 cache_values = [input_tensor.layout, chunks, dim, input_shape] 

59 return local_args, local_kwargs, cache_values 

60 

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

62 """ 

63 Infer output layouts for ChunkView operator. 

64 

65 Rules: 

66 1. Input must not have Partial status. 

67 2. Split dimension cannot be sharded (including StridedShard multi-axis mappings). 

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

69 4. Default: dim = 0 if not specified. 

70 5. Output count may be less than chunks if dimension size < chunks. 

71 6. All output layouts are identical to the input layout. 

72 

73 Args: 

74 cache_values (list): [input_layout, chunks, dim, input_shape] 

75 

76 Returns: 

77 tuple: ((output_layout_1, output_layout_2, ...), None) 

78 

79 Raises: 

80 ValueError: If any rule above is violated. 

81 TypeError: If chunks or dim is not an integer. 

82 """ 

83 input_layout = cache_values[0] 

84 chunks = cache_values[1] 

85 dim = cache_values[2] 

86 input_shape = cache_values[3] 

87 

88 if input_layout is None: 

89 raise ValueError( 

90 f"For {self.op_name}, input layout should not be None" 

91 ) 

92 

93 if not self._allow_partial_inputs: 

94 self._check_partial_inputs([input_layout]) 

95 

96 if not isinstance(chunks, int): 

97 raise TypeError( 

98 f"For {self.op_name}, chunks must be an integer, but got {type(chunks)}" 

99 ) 

100 if chunks < 1: 

101 raise ValueError( 

102 f"For {self.op_name}, chunks must be greater than 0, but got {chunks}" 

103 ) 

104 if not isinstance(dim, int): 

105 raise TypeError( 

106 f"For {self.op_name}, dim must be an integer, but got {type(dim)}" 

107 ) 

108 

109 alias_map = input_layout.alias_tensor_map 

110 ndim = len(alias_map) 

111 

112 original_dim = dim 

113 if dim < 0: 

114 dim = ndim + dim 

115 

116 if not 0 <= dim < ndim: 

117 raise ValueError( 

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

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

120 ) 

121 

122 mapping = alias_map[dim] 

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

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

125 else: 

126 is_sharded = mapping != "None" 

127 

128 if is_sharded: 

129 raise ValueError( 

130 f"For {self.op_name}, cannot split tensor at sharded axis[{dim}], " 

131 f"layout: {input_layout}" 

132 ) 

133 

134 output_num = self._calculate_output_count(input_shape[dim], chunks) 

135 

136 output_layouts = (input_layout,) * output_num 

137 return (output_layouts, None) 

138