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

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

16Distributed implementation for Slice operator. 

17""" 

18# pylint: disable=E0402 

19from typing import Callable, Optional, Tuple 

20 

21from .parallel_ops import DistributedOp 

22 

23 

24def _normalize_slice_args(x, begin, end): 

25 return (x, begin, end), {} 

26 

27 

28class SliceDistributedOp(DistributedOp): 

29 """Distributed implementation for Slice operator.""" 

30 

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

32 """ 

33 Preprocess arguments for Slice operator. 

34 

35 Args: 

36 args (tuple): Input arguments containing x, begin and end. 

37 kwargs (dict): Keyword arguments. 

38 

39 Returns: 

40 tuple: (local_args, local_kwargs, cache_values) 

41 """ 

42 args, _ = _normalize_slice_args(*args, **kwargs) 

43 input_tensor, begin, end = args 

44 local_args = (input_tensor.to_local(), begin, end) 

45 local_kwargs = {} 

46 cache_values = [input_tensor.layout, begin, end, input_tensor.shape] 

47 return local_args, local_kwargs, cache_values 

48 

49 def _is_shard_dim(self, layout): 

50 """return the shard num in each dim""" 

51 shard_dim = [] 

52 for axis_name in layout.alias_tensor_map: 

53 if axis_name == "None": 

54 shard_dim.append(1) 

55 continue 

56 if isinstance(axis_name, (tuple, list)): 

57 shard_num = 1 

58 for axis in axis_name: 

59 if axis != "None": 

60 shard_num *= layout.mesh.get_device_num_along_axis(axis) 

61 shard_dim.append(shard_num) 

62 continue 

63 shard_dim.append(layout.mesh.get_device_num_along_axis(axis_name)) 

64 return shard_dim 

65 

66 def _check_layout(self, layout, begin, end, shape): 

67 """check whether layout is valid""" 

68 if len(layout) != 1: 

69 raise ValueError(f"Layout must be a tuple of length 1, but got {len(layout)}") 

70 layout = layout[0] 

71 shard_dim = self._is_shard_dim(layout) 

72 for i, _ in enumerate(begin): 

73 if (shard_dim[i] != 1 and end[i] - begin[i] != shape[i]) and shape[i] != -1: 

74 raise ValueError( 

75 f"Slice: When a dimension({i}) is not fully fetched, the dimension can not be split now, " 

76 f"the begin is {begin}, the end is {end}, the shape is {shape}, layout is {layout.to_dict()}") 

77 return shard_dim 

78 

79 def infer_layout(self, cache_values: list) -> Tuple[tuple, tuple]: # pylint: disable=W0221 

80 """ 

81 Infer output layout for Slice operator. 

82 

83 Rules: 

84 1. Input must not have Partial status. 

85 2. begin, end and global_shape must have the same rank as input layout. 

86 3. Any sharded dimension must be fully fetched. 

87 4. Output layout is identical to the input layout. 

88 

89 Args: 

90 cache_values (list): [input_layout, begin, end, global_shape]. 

91 

92 Returns: 

93 tuple: ((output_layout,), (new_begin, new_end)). 

94 

95 Raises: 

96 ValueError: If input has Partial status, arguments rank mismatch, or a sharded 

97 dimension is not fully fetched. 

98 """ 

99 layout, begin, end, global_shape = cache_values 

100 self._check_partial_inputs([layout]) 

101 

102 if len(begin) != len(end) or len(begin) != len(global_shape): 

103 raise ValueError( 

104 f"For {self.op_name}, begin, end and global_shape must have the same length, " 

105 f"but got begin: {len(begin)}, end: {len(end)}, global_shape: {len(global_shape)}" 

106 ) 

107 if len(begin) != len(layout.alias_tensor_map): 

108 raise ValueError( 

109 f"For {self.op_name}, slice arguments rank must match input layout rank, " 

110 f"but got args rank: {len(begin)} and layout rank: {len(layout.alias_tensor_map)}" 

111 ) 

112 

113 shard_dim = self._check_layout((layout,), begin, end, global_shape) 

114 new_begin = tuple(begin[i] // shard_dim[i] for i in range(len(begin))) 

115 new_end = tuple(end[i] // shard_dim[i] for i in range(len(end))) 

116 return ((layout,), (new_begin, new_end)) 

117 

118 def get_expand_impl(self, func: Optional[Callable], infer_result: tuple, # pylint: disable=W0221 

119 cache_values: list) -> Optional[Callable]: 

120 """ 

121 Return a custom Slice implementation when local begin/end need adjustment. 

122 

123 Args: 

124 func: Original operator callable. 

125 infer_result (tuple): ((output_layout,), (new_begin, new_end)) from infer_layout. 

126 cache_values (list): [input_layout, begin, end, global_shape]. 

127 

128 Returns: 

129 callable | None: expand_impl closure when local slice bounds differ, else None. 

130 """ 

131 if func is None: 

132 return None 

133 

134 begin = cache_values[1] 

135 end = cache_values[2] 

136 new_begin, new_end = infer_result[1] 

137 if begin == new_begin and end == new_end: 

138 return None 

139 

140 def expand_impl(input_tensor: object, *_unused_args: object) -> object: 

141 """Call Slice with local slice bounds.""" 

142 return func(input_tensor, new_begin, new_end) 

143 

144 return expand_impl