Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / shard / ops / parallel_elementwise.py: 87%
255 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-07-06 05:41 +0800
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 Element-wise operator.
17"""
19import copy
20from typing import Callable, Optional, Tuple
22from .parallel_ops import DistributedOp
25def _unwrap_local_value(value):
26 """Convert DTensor-like values to local tensors while preserving containers."""
27 if hasattr(value, "_layout"):
28 return value.to_local()
29 if isinstance(value, tuple):
30 return tuple(_unwrap_local_value(item) for item in value)
31 if isinstance(value, list):
32 return [_unwrap_local_value(item) for item in value]
33 return value
36def _collect_layout_and_shape(value):
37 """Collect layout and shape from one argument for layout inference cache."""
38 layout = value.layout if hasattr(value, "_layout") else None
39 shape = value.shape if hasattr(value, "shape") else None
40 return layout, shape
43def _build_elementwise_cache_values(args, kwargs):
44 """Build cache_values from the real element-wise input arguments."""
45 input_layouts = []
46 input_shapes = []
47 for value in args:
48 layout, shape = _collect_layout_and_shape(value)
49 input_layouts.append(layout)
50 input_shapes.append(shape)
51 for value in kwargs.values():
52 layout, shape = _collect_layout_and_shape(value)
53 input_layouts.append(layout)
54 input_shapes.append(shape)
55 return [*input_layouts, input_shapes]
58class ElementWiseDistributedOp(DistributedOp):
59 """
60 Base class for distributed element-wise operators.
62 Supports broadcasting following broadcasting rules and handles
63 distributed tensor layouts with proper sharding strategy inference.
65 Args:
66 op_name (str): Name of the operator to register.
67 """
69 def preprocess(self, args: tuple, kwargs: dict) -> tuple:
70 """
71 Preprocess arguments for element-wise operators.
73 NOTE: aclop packed-args normalization (for MindSpore aclop operators
74 like Mod, StopGradient that pack args as ``(prim, name, (real_args...))``)
75 is handled upstream in ``OpDispatcher._dispatch_layout_infer`` via
76 ``_normalize_aclop_args``. This method receives clean unpacked args.
78 Args:
79 args (tuple): Positional arguments passed to the operator.
80 kwargs (dict): Keyword arguments passed to the operator.
82 Returns:
83 tuple: (local_args, local_kwargs, cache_values)
84 """
85 local_kwargs = {key: _unwrap_local_value(value) for key, value in kwargs.items()}
87 local_args = tuple(_unwrap_local_value(arg) for arg in args)
88 cache_values = _build_elementwise_cache_values(args, kwargs)
89 return local_args, local_kwargs, cache_values
91 def infer_layout(self, cache_values: list) -> Tuple[tuple, None]: # pylint: disable=W0221
92 """
93 Infer output layouts for element-wise operations with broadcasting support.
95 Rules:
96 1. Inputs must not have Partial status unless the operator explicitly allows it.
97 2. Input shapes must be broadcast-compatible when shape information is available.
98 3. Broadcasting dimensions cannot be sharded.
99 4. Non-broadcast sharding patterns must be compatible across inputs.
100 5. Output layout uses the merged sharding strategy and merged Partial status.
102 Args:
103 cache_values (list): [input_layout_0, ..., input_layout_n, input_shapes].
105 Returns:
106 tuple: ((output_layout,), None)
108 Raises:
109 ValueError: If input layouts are not compatible for broadcasting.
110 """
111 layouts = tuple(cache_values[:-1])
112 input_shapes = cache_values[-1] if cache_values else None
114 if not layouts:
115 return None
117 valid_layouts = [layout for layout in layouts if layout is not None]
119 if not valid_layouts:
120 return None
122 if not self._allow_partial_inputs:
123 self._check_partial_inputs(layouts)
125 if len(valid_layouts) == 1:
126 return ((copy.deepcopy(valid_layouts[0]),), None)
128 if not input_shapes:
129 return ((self._handle_no_input_shapes(valid_layouts),), None)
131 aligned_layouts, aligned_shapes = self._align_layouts_and_shapes(layouts, input_shapes)
133 if len(aligned_layouts) <= 1 or len(aligned_layouts) != len(aligned_shapes):
134 return ((copy.deepcopy(valid_layouts[0]),), None)
136 output_shape = self._compute_output_shape(aligned_shapes)
137 merged_tensor_map, merged_partial = self._merge_all_layouts(
138 aligned_layouts,
139 aligned_shapes,
140 output_shape,
141 layouts
142 )
144 self._check_all_inputs_broadcasts_and_partial(aligned_layouts, aligned_shapes, output_shape)
146 output_layout = self._create_output_layout(aligned_layouts[0], merged_tensor_map, merged_partial)
147 return ((output_layout,), None)
149 def _handle_no_input_shapes(self, valid_layouts):
150 """
151 Handle the case when input shapes are not available.
152 """
153 first_layout = valid_layouts[0]
154 first_alias_map = first_layout.alias_tensor_map
155 for layout in valid_layouts[1:]:
156 if layout.alias_tensor_map != first_alias_map:
157 raise ValueError(
158 f"For {self.op_name}, cannot infer layout without shapes: "
159 f"mismatched alias_tensor_map {first_alias_map} vs {layout.alias_tensor_map}."
160 )
161 return copy.deepcopy(first_layout)
163 def _align_layouts_and_shapes(self, layouts, input_shapes):
164 """
165 Align layouts with shapes by position, skipping None layouts.
166 """
167 aligned_layouts = []
168 aligned_shapes = []
169 for layout, shape in zip(layouts, input_shapes):
170 if layout is None:
171 continue
172 aligned_layouts.append(layout)
173 aligned_shapes.append(shape)
174 return aligned_layouts, aligned_shapes
176 def _compute_output_shape(self, aligned_shapes):
177 """
178 Compute broadcasted output shape from all input shapes.
179 """
180 output_shape = aligned_shapes[0]
181 for shape in aligned_shapes[1:]:
182 output_shape = self._broadcast_shapes(output_shape, shape)
183 return output_shape
185 def _merge_all_layouts(self, aligned_layouts, aligned_shapes, output_shape, layouts):
186 """
187 Merge all input layouts sequentially to get final tensor_map and partial status.
188 """
189 base_layout = aligned_layouts[0]
191 merged_tensor_map = self._merge_tensor_maps_for_broadcast(
192 aligned_layouts[0],
193 aligned_layouts[1],
194 aligned_shapes[0],
195 aligned_shapes[1],
196 output_shape
197 )
199 merged_partial = self._merge_partial_status(
200 base_layout.partial,
201 aligned_layouts[1].partial,
202 merged_tensor_map,
203 aligned_layouts[0].tensor_map if aligned_layouts[0].tensor_map else tuple(),
204 aligned_layouts[1].tensor_map if aligned_layouts[1].tensor_map else tuple(),
205 layouts
206 )
208 for i in range(2, len(aligned_layouts)):
209 temp_layout = self._create_output_layout(base_layout, merged_tensor_map, merged_partial)
210 merged_tensor_map = self._merge_tensor_maps_for_broadcast(
211 temp_layout,
212 aligned_layouts[i],
213 output_shape,
214 aligned_shapes[i],
215 output_shape
216 )
217 merged_partial = self._merge_partial_status(
218 merged_partial,
219 aligned_layouts[i].partial,
220 merged_tensor_map,
221 temp_layout.tensor_map if temp_layout.tensor_map else tuple(),
222 aligned_layouts[i].tensor_map if aligned_layouts[i].tensor_map else tuple(),
223 layouts
224 )
226 return merged_tensor_map, merged_partial
228 def _merge_partial_status(self, partial1, partial2, merged_tensor_map, tensor_map1, tensor_map2, layouts):
229 """
230 Merge partial status from two inputs.
232 Rules:
233 1. Both None → None
234 2. One None → Use the other
235 3. Both not None and same → Use it
236 4. Both not None and different → Error
237 5. Check Shard + Partial conflicts for each input
239 Args:
240 partial1: Partial status list from first input
241 partial2: Partial status list from second input
242 merged_tensor_map: Merged tensor map for output
243 tensor_map1: Tensor map of first input
244 tensor_map2: Tensor map of second input
246 Returns:
247 List: Merged partial status
249 Raises:
250 ValueError: If partial operations conflict or Shard+Partial conflict found
251 """
252 # Check Shard + Partial conflicts for input1
253 self._check_shard_partial_conflict(tensor_map1, partial1, layouts)
255 # Check Shard + Partial conflicts for input2
256 self._check_shard_partial_conflict(tensor_map2, partial2, layouts)
258 # Determine mesh dimension from partial lists
259 mesh_dim = max(len(partial1) if partial1 else 0, len(partial2) if partial2 else 0)
261 merged_partial = [None] * mesh_dim
263 for i in range(mesh_dim):
264 op1 = partial1[i] if partial1 and i < len(partial1) else None
265 op2 = partial2[i] if partial2 and i < len(partial2) else None
267 # Both have partial status with different operations
268 if op1 is not None and op2 is not None and op1 != op2:
269 raise ValueError(
270 f"For {self.op_name}, partial operations should be same for device axis {i}, "
271 f"but got {op1} and {op2}"
272 )
274 # Merge: prefer non-None, or either if both same
275 if op1 is not None:
276 merged_partial[i] = op1
277 elif op2 is not None:
278 merged_partial[i] = op2
280 # Check final output for Shard + Partial conflicts
281 self._check_shard_partial_conflict(merged_tensor_map, merged_partial, layouts)
283 return merged_partial
285 def _check_shard_partial_conflict(self, tensor_map, partial_list, layouts):
286 """
287 Check for conflicts between Shard and Partial on same device axis.
289 Args:
290 tensor_map: Tensor map to check
291 partial_list: Partial status list
293 Raises:
294 ValueError: If Shard and Partial conflict found
295 """
296 if not partial_list:
297 return
299 mesh_dim = len(partial_list)
301 # Collect all device axis used for sharding
302 sharded_axis = set()
303 if tensor_map:
304 for map_val in tensor_map:
305 if isinstance(map_val, tuple):
306 for sub_val in map_val:
307 if sub_val != -1:
308 # Convert to device axis index
309 axis_idx = mesh_dim - 1 - sub_val
310 sharded_axis.add(axis_idx)
311 elif map_val != -1:
312 axis_idx = mesh_dim - 1 - map_val
313 sharded_axis.add(axis_idx)
315 # Check if any sharded axis has partial status
316 for axis_idx in sharded_axis:
317 if 0 <= axis_idx < len(partial_list) and partial_list[axis_idx] is not None:
318 raise ValueError(
319 f"For {self.op_name}, Shard and Partial should not coexist on same device axis "
320 f"{axis_idx}, but got Partial({partial_list[axis_idx]}). "
321 f"Please check layouts: {layouts}."
322 )
324 def _check_all_inputs_broadcasts_and_partial(self, layouts, input_shapes, output_shape):
325 """
326 Check if any input broadcasts and has Partial status.
327 """
328 for i, (layout, input_shape) in enumerate(zip(layouts, input_shapes)):
329 if layout is None:
330 continue
332 input_name = f"input{i+1}"
334 input_len = len(input_shape)
335 output_len = len(output_shape)
337 if input_len < output_len:
338 aligned_input_shape = (1,) * (output_len - input_len) + tuple(input_shape)
339 else:
340 aligned_input_shape = input_shape
342 broadcasts = False
343 for in_dim, out_dim in zip(aligned_input_shape, output_shape):
344 if in_dim == 1 and out_dim > 1:
345 broadcasts = True
346 break
348 if broadcasts and layout.is_partial():
349 raise ValueError(
350 f"For {self.op_name}, {input_name} has Partial status and broadcasts. "
351 f"Should be without Partial status for broadcasting without communication"
352 )
354 def _broadcast_shapes(self, shape1, shape2):
355 """
356 Calculate the broadcasted shape of two shapes according to broadcasting rules.
358 Broadcasting rules:
359 1. If two arrays have different numbers of dimensions, pad the shape of the
360 lower-dimensional array with 1s on the left until both shapes have the same length.
361 2. If two arrays have the same number of dimensions but different lengths in some
362 dimensions, dimensions with length 1 will be expanded to match the other array's
363 dimension length.
364 3. If two arrays have the same number of dimensions but any dimension has different
365 lengths and neither is 1, raise an error.
367 Args:
368 shape1 (tuple): Shape of the first tensor, e.g., (3, 1, 5)
369 shape2 (tuple): Shape of the second tensor, e.g., (4, 5)
371 Returns:
372 tuple: Broadcasted shape, e.g., (3, 4, 5)
374 Raises:
375 ValueError: If shapes cannot be broadcast together.
376 """
377 # Rule 1: Right-align, pad with 1s on the left to make dimensions equal
378 len1, len2 = len(shape1), len(shape2)
379 max_len = max(len1, len2)
381 padded_shape1 = (1,) * (max_len - len1) + tuple(shape1)
382 padded_shape2 = (1,) * (max_len - len2) + tuple(shape2)
384 # Rules 2 and 3: Check if each dimension can be broadcast
385 result_shape = []
386 for dim1, dim2 in zip(padded_shape1, padded_shape2):
387 if dim1 == dim2:
388 # Dimensions are the same, use directly
389 result_shape.append(dim1)
390 elif dim1 == 1:
391 # First shape has 1 in this dimension, expand to dim2
392 result_shape.append(dim2)
393 elif dim2 == 1:
394 # Second shape has 1 in this dimension, expand to dim1
395 result_shape.append(dim1)
396 else:
397 # Rule 3: Dimensions are different and neither is 1, cannot broadcast
398 raise ValueError(
399 f"For {self.op_name}, shapes {shape1} and {shape2} cannot be broadcast together. "
400 f"Dimension mismatch: {dim1} vs {dim2}"
401 )
403 return tuple(result_shape)
405 def _align_tensor_maps_for_broadcast(self, layout1, layout2, shape1, shape2):
406 """
407 Align tensor_maps of two layouts to support broadcasting.
409 When two tensors have different dimensions, the tensor_map of the
410 lower-dimensional tensor is padded with -1 (indicating no sharding) at the front.
412 Args:
413 layout1: Layout of the first tensor
414 layout2: Layout of the second tensor
415 shape1 (tuple): Global shape of the first tensor
416 shape2 (tuple): Global shape of the second tensor
418 Returns:
419 tuple: (aligned_map1, aligned_map2) - Aligned tensor_maps
420 """
421 len1, len2 = len(shape1), len(shape2)
422 max_len = max(len1, len2)
424 map1 = layout1.tensor_map if layout1.tensor_map else tuple([-1] * len1)
425 map2 = layout2.tensor_map if layout2.tensor_map else tuple([-1] * len2)
427 aligned_map1 = (-1,) * (max_len - len1) + map1
428 aligned_map2 = (-1,) * (max_len - len2) + map2
430 return aligned_map1, aligned_map2
432 def _normalize_tensor_map_element(self, map_element):
433 """
434 Normalize a tensor_map element to a tuple of device axis for unified processing.
436 Args:
437 map_element: Element from tensor_map, can be:
438 - int: -1 (no sharding) or device axis index
439 - tuple: multiple device axis
441 Returns:
442 tuple: Tuple of device axis (empty tuple if not sharded)
443 """
444 if map_element == -1:
445 return ()
446 if isinstance(map_element, int):
447 return (map_element,)
448 if isinstance(map_element, tuple):
449 return tuple(dim for dim in map_element if dim != -1)
450 return ()
452 def _denormalize_tensor_map_element(self, device_axis_tuple):
453 """
454 Convert a tuple of device axis back to tensor_map element format.
456 Args:
457 device_axis_tuple (tuple): Tuple of device axis
459 Returns:
460 int or tuple: -1 if empty, single int if one element, tuple if multiple elements
461 """
462 if not device_axis_tuple:
463 return -1
464 if len(device_axis_tuple) == 1:
465 return device_axis_tuple[0]
466 return device_axis_tuple
468 def _merge_tensor_maps_for_broadcast(self, layout1, layout2, shape1, shape2, output_shape):
469 """
470 Merge tensor_maps of two inputs to generate output tensor_map.
472 This method handles both simple int-type and complex tuple-type tensor_map elements,
473 ensuring correct sharding strategy for the broadcasted output.
475 Args:
476 layout1: Layout of the first input
477 layout2: Layout of the second input
478 shape1 (tuple): Global shape of the first input
479 shape2 (tuple): Global shape of the second input
480 output_shape (tuple): Global shape of the output
482 Returns:
483 tuple: Merged tensor_map for the output
485 Raises:
486 ValueError: If sharding strategies conflict or broadcasting dimension is sharded
487 """
488 map1, map2 = self._align_tensor_maps_for_broadcast(layout1, layout2, shape1, shape2)
490 len1, len2 = len(shape1), len(shape2)
491 max_len = len(output_shape)
492 padded_shape1 = (1,) * (max_len - len1) + tuple(shape1)
493 padded_shape2 = (1,) * (max_len - len2) + tuple(shape2)
495 merged_map = []
496 for i, (dim1, dim2, out_dim) in enumerate(zip(padded_shape1, padded_shape2, output_shape)):
497 m1, m2 = map1[i], map2[i]
499 m1_axis = self._normalize_tensor_map_element(m1)
500 m2_axis = self._normalize_tensor_map_element(m2)
502 m1_axis_for_compare = frozenset(m1_axis)
503 m2_axis_for_compare = frozenset(m2_axis)
505 m1_is_sharded = bool(m1_axis)
506 m2_is_sharded = bool(m2_axis)
508 if not m1_is_sharded and not m2_is_sharded:
509 merged_map.append(-1)
511 elif not m1_is_sharded:
512 if dim2 == 1 and out_dim > 1:
513 raise ValueError(
514 f"For {self.op_name}, dimension {i} of second input has size 1 "
515 f"but is sharded on device axis {m2_axis}. "
516 f"Broadcasting dimension cannot be sharded."
517 )
518 merged_map.append(self._denormalize_tensor_map_element(m2_axis))
520 elif not m2_is_sharded:
521 if dim1 == 1 and out_dim > 1:
522 raise ValueError(
523 f"For {self.op_name}, dimension {i} of first input has size 1 "
524 f"but is sharded on device axis {m1_axis}. "
525 f"Broadcasting dimension cannot be sharded."
526 )
527 merged_map.append(self._denormalize_tensor_map_element(m1_axis))
529 else:
530 if m1_axis_for_compare != m2_axis_for_compare:
531 raise ValueError(
532 f"For {self.op_name}, inputs should have same sharding pattern, "
533 f"but got confilcting sharding at dimension {i}, "
534 f"input1 shaded on {m1_axis} and input2 shaded on {m2_axis}."
535 )
537 if (dim1 == 1 or dim2 == 1) and dim1 != dim2:
538 raise ValueError(
539 f"For {self.op_name}, dimension {i} is broadcast from size 1 "
540 f"to {out_dim} but is sharded on device axis {m1_axis}. "
541 f"Broadcasting dimension cannot be sharded."
542 )
544 merged_map.append(self._denormalize_tensor_map_element(m1_axis))
546 return tuple(merged_map)
548 def _create_output_layout(self, base_layout, output_tensor_map, partial_list=None):
549 """
550 Create output layout based on input layout.
552 Args:
553 base_layout: Base layout (usually from the first input)
554 output_tensor_map (tuple): Tensor_map for the output
555 partial_list (list): Partial status list for the output
557 Returns:
558 Layout: New Layout object with updated tensor_map and alias_tensor_map
559 """
560 new_layout = copy.deepcopy(base_layout)
561 new_layout.set_tensor_map(output_tensor_map)
563 alias_tensor_map = []
564 for tensor_dim in output_tensor_map:
565 if tensor_dim == -1:
566 alias_tensor_map.append("None")
567 elif isinstance(tensor_dim, tuple):
568 alias_tuple = tuple(
569 base_layout.alias_name[len(base_layout.alias_name) - 1 - dim]
570 for dim in tensor_dim
571 if dim != -1
572 )
573 alias_tensor_map.append(alias_tuple if alias_tuple else "None")
574 else:
575 alias_tensor_map.append(
576 base_layout.alias_name[len(base_layout.alias_name) - 1 - tensor_dim]
577 )
579 new_layout.set_alias_tensor_map(tuple(alias_tensor_map))
581 # Set partial status if provided
582 if partial_list:
583 for i, partial_op in enumerate(partial_list):
584 if partial_op is not None and i < len(new_layout.alias_name):
585 new_layout.set_partial_by_dev_axis(new_layout.alias_name[i], partial_op)
587 return new_layout
590class ElementWiseWithPartialDistributedOp(ElementWiseDistributedOp):
591 """
592 Base class for elementwise operations that support partial status propagation.
593 """
594 def __init__(self, op_name):
595 super().__init__(op_name)
596 self._allow_partial_inputs = True
599class AddDistributedOp(ElementWiseWithPartialDistributedOp):
600 """
601 Distributed implementation for Add operator.
603 This operator supports partial status propagation from inputs to output,
604 which is useful for operations like gradient accumulation where partial
605 results need to be preserved through the computation graph.
606 """
607 def infer_layout(self, cache_values: list) -> Tuple[tuple, None]: # pylint: disable=W0221
608 """
609 Infer output layout for Add operator.
611 Rules:
612 1. Follow element-wise broadcasting and sharding merge rules.
613 2. Partial status is allowed and propagated.
614 3. Propagated Partial status must be "sum" or None.
616 Args:
617 cache_values (list): [input_layout_0, ..., input_layout_n, input_shapes].
619 Returns:
620 tuple: ((output_layout,), None)
622 Raises:
623 ValueError: If propagated Partial status is not "sum" or None.
624 """
625 infer_result = super().infer_layout(cache_values)
626 if infer_result is None:
627 return infer_result
629 output_layout = infer_result[0][0]
630 for i, partial_type in enumerate(output_layout.partial):
631 if partial_type is not None and partial_type != "sum":
632 raise ValueError(
633 f"For {self.op_name}, inputs partial status should be 'sum' or None, "
634 f"but got {partial_type} at index {i}."
635 )
637 return infer_result
639 def get_expand_impl(self, func: Optional[Callable], infer_result: tuple, # pylint: disable=W0221
640 cache_values: list) -> Optional[Callable]:
641 """
642 Get expand implementation for the operator
643 """
644 layouts = tuple(cache_values[:-1])
645 x1_layout = layouts[0]
646 x2_layout = layouts[1]
647 x1_partial = x1_layout.is_partial() if x1_layout is not None else False
648 x2_partial = x2_layout.is_partial() if x2_layout is not None else False
650 if x1_partial == x2_partial:
651 return None
653 output_layout = infer_result[0][0]
654 scaling_factor = 1
655 for i, partial_type in enumerate(output_layout.partial):
656 if partial_type == "sum":
657 scaling_factor *= output_layout.mesh_shape[i]
659 # use expand_impl only when one of x1 and x2 is with partial placement.
660 def _expand_impl1(x1, x2, *args, **kwargs):
661 add_out = func(x1 / scaling_factor, x2, *args, **kwargs)
662 return add_out
664 def _expand_impl2(x1, x2, *args, **kwargs):
665 add_out = func(x1, x2 / scaling_factor, *args, **kwargs)
666 return add_out
668 return _expand_impl1 if not x1_partial else _expand_impl2