Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / shard / ops / parallel_reshape.py: 83%
182 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-07-06 05:41 +0800
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 Reshape operator.
17"""
19from typing import Callable, Optional, Tuple
21from hyper_parallel.core.dtensor.layout import Layout
22from hyper_parallel.platform import get_platform
23from .parallel_ops import DistributedOp
24platform = get_platform()
25Tensor = platform.Tensor
28def _normalize_reshape_args(x, *shape, **kwargs):
29 """Normalize reshape/view arguments into positional args and empty kwargs."""
30 unexpected_kwargs = set(kwargs) - {'shape'}
31 if unexpected_kwargs:
32 unexpected = next(iter(unexpected_kwargs))
33 raise TypeError(f"reshape got an unexpected keyword argument '{unexpected}'.")
34 if shape and 'shape' in kwargs:
35 raise TypeError("reshape got shape from both args and kwargs.")
36 if not shape and 'shape' in kwargs:
37 shape = (kwargs['shape'],)
38 if not shape:
39 raise TypeError("reshape missing required shape argument.")
40 return (x,) + shape, {}
43def _filter_none_split_tensor_map(tensor_map, mesh_shape):
44 """
45 Filter out the elements in tensor_map where the size of the corresponding dimension in device_matrix is 1.
47 Args:
48 tensor_map (list): A list of tensor mappings, which may contain integers or tuples.
49 device_matrix (list): A device matrix representing the device distribution across each dimension.
51 Returns:
52 list: The filtered list of tensor mappings, where invalid mappings are replaced with -1 or valid mappings are
53 retained.
54 """
55 filtered_tensor_map = []
56 for item in tensor_map:
57 if isinstance(item, tuple):
58 filtered = []
59 for i in item:
60 if mesh_shape[-1 - i] != 1:
61 filtered.append(i)
62 if len(filtered) == 0:
63 filtered_tensor_map.append(-1)
64 elif len(filtered) == 1:
65 filtered_tensor_map.append(filtered[0])
66 else:
67 filtered_tensor_map.append(tuple(filtered))
68 else:
69 filtered_tensor_map.append(item if mesh_shape[-1 - item] != 1 else -1)
70 return filtered_tensor_map
73class ReshapeDistributedOp(DistributedOp):
74 """Distributed implementation for Reshape operator."""
76 def __init__(self, op_name):
77 super().__init__(op_name)
78 self._allow_partial_inputs = True
80 def _get_dynamic_shape_info(self, shape):
81 total_size = 1
82 dynamic_axis = -1
83 for axis, s in enumerate(shape):
84 total_size *= s
85 if s < 0:
86 dynamic_axis = axis
87 return total_size < 0, dynamic_axis, total_size
89 def _handle_dynamic_shape(self, input_shape, output_shape):
90 """
91 Check dynamic shape. Calculate unknown axis if one of input and output shape is known. If both are unknown,
92 calculate the relative multiple.
93 [2, -1, 8], [4, -1, 8] -> [2, -2, 8], [4, -1, 8]
94 """
95 input_shape = list(input_shape)
96 output_shape = list(output_shape)
97 is_input_dynamic, input_dynamic_axis, input_total_size = self._get_dynamic_shape_info(input_shape)
98 is_output_dynamic, output_dynamic_axis, output_total_size = self._get_dynamic_shape_info(output_shape)
99 dynamic_can_shard = False
100 if not is_input_dynamic and not is_output_dynamic:
101 if input_total_size != output_total_size:
102 raise ValueError(f"The total elements number of input shape {input_shape} and output shape "
103 f"{output_shape} are different.")
104 return input_shape, output_shape, dynamic_can_shard
106 if not is_input_dynamic:
107 accurate_output_shape = output_shape
108 accurate_output_shape[output_dynamic_axis] = -input_total_size // output_total_size
109 return input_shape, accurate_output_shape, dynamic_can_shard
111 if not is_output_dynamic:
112 accurate_input_shape = input_shape
113 accurate_input_shape[input_dynamic_axis] = -output_total_size // input_total_size
114 return accurate_input_shape, output_shape, dynamic_can_shard
116 if output_total_size >= input_total_size:
117 output_shape[output_dynamic_axis] = -(input_total_size // output_total_size)
118 dynamic_can_shard = True
119 else:
120 input_shape[input_dynamic_axis] = -(output_total_size // input_total_size)
121 return input_shape, output_shape, dynamic_can_shard
123 def _merge_unshared_axis(self, global_shape, tensor_map):
124 """
125 Merge those axes that are not sharded to the high dimension which is shared.
126 shape[4, 2, 6, 8], tensor map[-1, -1, 0, -1] -> merged shape[8, 48]
128 Returns:
129 tuple: (merged_shape, merge_tensor_map).
130 merge_tensor_map may contain -1 for merged unsharded axis groups.
131 """
132 merged_size = 1
133 merged_shape = []
134 merged_tensor_map = []
135 for axis in range(len(global_shape) - 1, -1, -1):
136 merged_size *= global_shape[axis]
137 if tensor_map[axis] != -1:
138 merged_shape.insert(0, merged_size)
139 merged_tensor_map.insert(0, tensor_map[axis])
140 merged_size = 1
141 if tensor_map[0] == -1:
142 merged_shape.insert(0, merged_size)
143 merged_tensor_map.insert(0, -1)
144 return merged_shape, merged_tensor_map
147 def _cal_output_layout_and_dst_shape(self, output_tensor_map, dst_shape, x_dict):
148 """
149 calculate output layout tensor map and local dst shape.
150 """
151 x_mesh_shape = x_dict["mesh_shape"]
152 output_map = []
153 local_dst_shape = []
154 for idx, map_id in enumerate(output_tensor_map):
155 if isinstance(map_id, tuple):
156 shard_size = 1
157 map_idx = []
158 for shard_id in map_id:
159 map_idx.append(x_dict["alias_name"][-1 - shard_id])
160 shard_size *= x_mesh_shape[-1 - shard_id]
161 output_map.append(tuple(map_idx))
162 local_dst_shape.append(dst_shape[idx] // shard_size if dst_shape[idx] > 0 else -1)
163 continue
164 if map_id < 0:
165 output_map.append("None")
166 local_dst_shape.append(dst_shape[idx] if dst_shape[idx] > 0 else -1)
167 else:
168 output_map.append(x_dict["alias_name"][-1 - map_id])
169 local_dst_shape.append(dst_shape[idx] // x_mesh_shape[-1 - map_id] if dst_shape[idx] > 0 else -1)
170 return output_map, local_dst_shape
172 def _normalize_shape(self, dst_shape):
173 """Normalize dst_shape to list format."""
174 if isinstance(dst_shape, Tensor):
175 dst_shape = dst_shape.tolist()
176 if not isinstance(dst_shape, (list, tuple)):
177 raise ValueError("Shape should be a tensor or a tuple or a list.")
178 return dst_shape
180 def _compute_output_tensor_map(self, merged_shape, merge_tensor_map, dst_shape, x_mesh_shape, dynamic_can_shard,
181 input_shape, x_map):
182 """Compute output tensor_map from merged information.
184 Args:
185 merged_shape: Merged shape from _merge_unshared_axis
186 merge_tensor_map: Merged tensor_map from _merge_unshared_axis
187 dst_shape: Target shape
188 x_mesh_shape: Mesh shape
189 dynamic_can_shard: Whether dynamic shape can be sharded
190 input_shape: Original input shape
191 x_map: Input tensor_map
193 Returns:
194 list: Output tensor_map
195 """
196 output_tensor_map = []
197 cur_axis = len(merged_shape) - 1
198 cur_size = merged_shape[cur_axis]
200 for shape in reversed(dst_shape):
201 if cur_size % shape != 0:
202 raise ValueError(f"Can not reshape {input_shape} to {dst_shape} with tensor map {x_map}")
203 cur_size = cur_size // shape
205 if cur_size == 1:
206 map_val = merge_tensor_map[cur_axis]
207 if map_val != -1:
208 self._validate_reshape_shard(
209 map_val, x_mesh_shape, shape,
210 dynamic_can_shard, input_shape, x_map, dst_shape
211 )
212 output_tensor_map.insert(0, map_val)
213 cur_axis -= 1
214 cur_size = merged_shape[cur_axis]
215 else:
216 output_tensor_map.insert(0, -1)
218 return output_tensor_map
220 def _validate_reshape_shard(self, map_val, x_mesh_shape, shape,
221 dynamic_can_shard, input_shape, x_map, dst_shape):
222 """Validate that a sharded axis can be reshaped to the target shape dimension."""
223 if isinstance(map_val, tuple):
224 shard_size = 1
225 for axis in map_val:
226 shard_size *= x_mesh_shape[-axis - 1]
227 else:
228 shard_size = x_mesh_shape[-map_val - 1]
230 if shape < 0:
231 if not dynamic_can_shard:
232 raise ValueError(f"Can not reshape {input_shape} to {dst_shape} with tensor map {x_map}")
233 elif shard_size > shape or shape % shard_size != 0:
234 raise ValueError(f"Can not reshape {input_shape} to {dst_shape} with tensor map {x_map}")
236 def _apply_partial_status(self, x_layout, out_layout):
237 """Apply partial status from input to output layout."""
238 if x_layout.is_partial():
239 input_partial = x_layout.partial
240 for i, partial_op in enumerate(input_partial):
241 if partial_op is not None and i < len(out_layout.alias_name):
242 out_layout.set_partial_by_dev_axis(out_layout.alias_name[i], partial_op)
244 def preprocess(self, args: tuple, kwargs: dict) -> tuple:
245 """
246 Preprocess arguments for Reshape operator.
248 Args:
249 args (tuple): Input tensor followed by target shape arguments.
250 kwargs (dict): Keyword arguments.
252 Returns:
253 tuple: (local_args, local_kwargs, cache_values)
254 """
255 args, _ = _normalize_reshape_args(*args, **kwargs)
256 input_tensor = args[0]
257 dst_shape = args[1:] if len(args) > 2 else args[1]
259 local_args = (input_tensor.to_local(), dst_shape)
260 local_kwargs = {}
261 cache_values = [input_tensor.layout, dst_shape, tuple(input_tensor.shape)]
262 return local_args, local_kwargs, cache_values
264 def _infer_reshape_layout(self, x_layout, dst_shape, input_shape):
265 """Infer reshape output layout and local destination shape."""
266 x_dict = x_layout.to_dict()
267 dst_shape = self._normalize_shape(dst_shape)
269 x_map = _filter_none_split_tensor_map(x_dict["tensor_map"], x_dict["mesh_shape"])
270 x_mesh_shape = x_dict["mesh_shape"]
272 input_shape, dst_shape, dynamic_can_shard = self._handle_dynamic_shape(input_shape, dst_shape)
273 merged_shape, merge_tensor_map = self._merge_unshared_axis(input_shape, x_map)
275 output_tensor_map = self._compute_output_tensor_map(
276 merged_shape, merge_tensor_map, dst_shape, x_mesh_shape, dynamic_can_shard, input_shape, x_map
277 )
279 output_layout = Layout(
280 mesh_shape=x_mesh_shape,
281 alias_name=x_layout.alias_name,
282 rank_list=x_layout.rank_list
283 )
284 output_map, local_dst_shape = self._cal_output_layout_and_dst_shape(output_tensor_map, dst_shape, x_dict)
285 out_layout = output_layout(*output_map)
287 self._apply_partial_status(x_layout, out_layout)
289 return out_layout, local_dst_shape
291 def infer_layout(self, cache_values: list) -> Tuple[tuple, list]: # pylint: disable=W0221
292 """
293 Infer output layout for Reshape operator.
295 Rules:
296 1. Partial input is allowed and preserved on the output layout.
297 2. Target shape must be a Tensor, tuple, or list.
298 3. Input and output total element counts must match after resolving one dynamic axis.
299 4. Reshape must preserve each device's local data slice; sharded axes can only be
300 split or merged when the shard boundary remains valid.
301 5. Output Partial status follows the input Partial status.
303 Args:
304 cache_values (list): [input_layout, dst_shape, input_shape].
306 Returns:
307 tuple: ((output_layout,), local_dst_shape)
309 Raises:
310 ValueError: If target shape is invalid or the reshape would change sharded slices.
311 """
312 if len(cache_values) != 3:
313 raise ValueError(
314 f"For {self.op_name}, cache_values length should be 3, but got {len(cache_values)}"
315 )
317 x_layout, dst_shape, input_shape = cache_values[0], cache_values[1], cache_values[2]
318 if x_layout is None:
319 raise ValueError(f"For {self.op_name}, reshape requires a valid input tensor layout.")
321 out_layout, local_dst_shape = self._infer_reshape_layout(x_layout, dst_shape, input_shape)
322 return ((out_layout,), local_dst_shape)
324 def get_expand_impl(self, func: Optional[Callable], infer_result: tuple, # pylint: disable=W0221
325 cache_values: list) -> Optional[Callable]:
326 """Return a closure that calls reshape/view with the inferred local target shape."""
327 del cache_values
328 if func is None:
329 return None
331 local_dst_shape = infer_result[1]
332 if local_dst_shape is None:
333 return None
335 def expand_impl(x: object, shape: object) -> object:
336 del shape
337 return func(x, local_dst_shape)
339 return expand_impl