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