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

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 Stack operator. 

17""" 

18 

19from typing import Tuple 

20 

21from hyper_parallel.core.dtensor.layout import Layout 

22from .parallel_ops import DistributedOp 

23 

24# pylint: disable=unused-argument 

25 

26 

27def _normalize_stack_args(tensors, dim=0, **kwargs): 

28 """ 

29 Normalize arguments for torch.stack. 

30 """ 

31 return (tensors,), {'dim': dim} 

32 

33 

34class StackDistributedOp(DistributedOp): 

35 """Distributed implementation for Stack operator.""" 

36 

37 def preprocess(self, args, kwargs): 

38 """ 

39 Preprocess input arguments and extract local components. 

40 

41 Normalizes args, explicitly extracts parameters, and prepares 

42 local tensors and cache values without validation logic. 

43 """ 

44 args, kwargs = _normalize_stack_args(*args, **kwargs) 

45 

46 # Explicit parameter extraction 

47 tensors = args[0] 

48 dim = kwargs['dim'] 

49 

50 # Extract local tensors and layouts 

51 local_tensors = tuple(t.to_local() if hasattr(t, "to_local") else t for t in tensors) 

52 layouts = [getattr(t, "layout", None) for t in tensors] 

53 

54 # Construct local args and kwargs for the inner op execution 

55 local_args = (local_tensors,) 

56 local_kwargs = {'dim': dim} 

57 

58 # Flatten layouts and append dim for caching and inference 

59 cache_values = layouts + [dim] 

60 

61 return local_args, local_kwargs, cache_values 

62 

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

64 """ 

65 Infer output layout for Stack operator. 

66 

67 Rules: 

68 1. At least one input DTensor is required. 

69 2. All input tensors must have identical layouts. 

70 3. A new Replicate dimension is inserted at the stack position. 

71 

72 Args: 

73 cache_values (list): [layout_0, ..., layout_n, dim] 

74 

75 Returns: 

76 tuple: ((output_layout,), None) 

77 

78 Raises: 

79 ValueError: If any rule above is violated. 

80 """ 

81 layouts = cache_values[:-1] 

82 dim = cache_values[-1] 

83 

84 if not layouts: 

85 raise ValueError( 

86 f"For {self.op_name}, stack requires at least one input tensor." 

87 ) 

88 

89 valid_layouts = [lyt for lyt in layouts if lyt is not None] 

90 

91 if not valid_layouts: 

92 raise ValueError( 

93 f"For {self.op_name}, stack requires at least one input DTensor." 

94 ) 

95 

96 if not self._allow_partial_inputs: 

97 self._check_partial_inputs(valid_layouts) 

98 

99 # Reference layout to validate consistency across all input tensors 

100 base_layout = valid_layouts[0] 

101 for layout in valid_layouts[1:]: 

102 if layout != base_layout: 

103 raise ValueError( 

104 f"For {self.op_name}, all input tensors must have the same layout, " 

105 f"but got base: {base_layout} and mismatched: {layout}" 

106 ) 

107 

108 ndim = len(base_layout.tensor_map) 

109 

110 # Normalize and validate the dimension. For stack, valid range is [-ndim - 1, ndim] 

111 actual_dim = dim if dim >= 0 else dim + ndim + 1 

112 if actual_dim < 0 or actual_dim > ndim: 

113 raise ValueError( 

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

115 f"expected to be in range of [{-ndim - 1}, {ndim}], but got {dim}" 

116 ) 

117 

118 # 2. Layout Inference Logic 

119 in_tensor_map = base_layout.tensor_map 

120 

121 # Insert an unsharded mapping (-1) at the newly created dimension 

122 output_tensor_map = list(in_tensor_map) 

123 output_tensor_map.insert(actual_dim, -1) 

124 

125 mesh_shape = base_layout.mesh_shape 

126 alias_name = base_layout.alias_name 

127 rank_list = base_layout.rank_list 

128 

129 def idx_to_alias(idx, aliases): 

130 """Map tensor_map index back to the alias string.""" 

131 if idx == -1: 

132 return "None" 

133 # Reverse indexing mapped to the framework's layout design 

134 return aliases[len(aliases) - idx - 1] 

135 

136 output_alias_map = tuple(idx_to_alias(idx, alias_name) for idx in output_tensor_map) 

137 

138 # Reconstruct the output layout 

139 output_layout = Layout( 

140 mesh_shape=mesh_shape, 

141 alias_name=alias_name, 

142 rank_list=rank_list 

143 ) 

144 

145 # Apply the placement mappings via the __call__ method 

146 output_layout = output_layout(*output_alias_map) 

147 

148 return ((output_layout,), None)