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

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 

19import copy 

20from typing import Tuple 

21 

22from .parallel_ops import DistributedOp 

23 

24 

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

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

27 

28 

29class ChunkViewDistributedOp(DistributedOp): 

30 """Distributed implementation for ChunkView operator.""" 

31 

32 @staticmethod 

33 def _calculate_output_count(dim_size, chunks): 

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

35 if dim_size == 0: 

36 return chunks 

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

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

39 return min(output_num, chunks) 

40 

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

42 """ 

43 Preprocess arguments for ChunkView operator. 

44 

45 Args: 

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

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

48 

49 Returns: 

50 tuple: (local_args, local_kwargs, cache_values) 

51 """ 

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

53 input_tensor, chunks, dim = args 

54 input_shape = input_tensor.shape 

55 

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

57 local_kwargs = {} 

58 

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

60 return local_args, local_kwargs, cache_values 

61 

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

63 """ 

64 Infer output layouts for ChunkView operator. 

65 

66 Rules: 

67 1. Input must not have Partial status. 

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

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

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

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

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

73 

74 Args: 

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

76 

77 Returns: 

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

79 

80 Raises: 

81 ValueError: If any rule above is violated. 

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

83 """ 

84 input_layout = cache_values[0] 

85 chunks = cache_values[1] 

86 dim = cache_values[2] 

87 input_shape = cache_values[3] 

88 

89 if input_layout is None: 

90 raise ValueError( 

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

92 ) 

93 

94 if not self._allow_partial_inputs: 

95 self._check_partial_inputs([input_layout]) 

96 

97 if not isinstance(chunks, int): 

98 raise TypeError( 

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

100 ) 

101 if chunks < 1: 

102 raise ValueError( 

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

104 ) 

105 if not isinstance(dim, int): 

106 raise TypeError( 

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

108 ) 

109 

110 alias_map = input_layout.alias_tensor_map 

111 ndim = len(alias_map) 

112 

113 original_dim = dim 

114 if dim < 0: 

115 dim = ndim + dim 

116 

117 if not 0 <= dim < ndim: 

118 raise ValueError( 

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

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

121 ) 

122 

123 mapping = alias_map[dim] 

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

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

126 else: 

127 is_sharded = mapping != "None" 

128 

129 if is_sharded: 

130 raise ValueError( 

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

132 f"layout: {input_layout}" 

133 ) 

134 

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

136 

137 output_layouts = tuple(copy.deepcopy(input_layout) for _ in range(output_num)) 

138 return (output_layouts, None)