Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / shard / ops / parallel_reduce.py: 89%
192 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 Reduce operator.
17"""
19from copy import deepcopy
20from typing import Sequence, Union, Tuple, List
22from hyper_parallel.core.dtensor.layout import Layout
23from hyper_parallel.platform import get_platform
24from .parallel_ops import DistributedOp
26platform = get_platform()
27Tensor = platform.Tensor
30StrOrTuple = Union[str, Tuple["StrOrTuple", ...], List["StrOrTuple"]]
33def _normalize_reduce_args(input_tensor, dim=None, keepdim=False, dtype=None):
34 """Normalize reduce-family arguments to consistent positional form.
36 Handles torch.sum / torch.mean / torch.prod / torch.all and MindSpore
37 SumExt / MeanExt / ReduceMax / MaxDim where *dim* and *keepdim* are
38 ordinary positional parameters. SumExt and MeanExt also pass a trailing
39 dtype slot positionally; normalize it here while keeping layout inference
40 based only on ``input_layout``, ``dim``, and ``keepdim``.
42 Args:
43 input_tensor: The input tensor (DTensor or Tensor).
44 dim: Dimension(s) to reduce. None means all dimensions, () means no reduction.
45 keepdim: Whether to retain reduced dimensions.
46 dtype: Optional output dtype.
48 Returns:
49 tuple: ``((input_tensor, dim, keepdim, dtype), {})`` — all-positional, empty kwargs.
50 """
51 return (input_tensor, dim, keepdim, dtype), {}
54class ReduceExtDistributedOpBase(DistributedOp):
55 """
56 Base class for distributed reduce operators.
58 Args:
59 op_name (str): Name of the operator to register.
60 partial_type (list): List of the operator for allreduce.
61 """
63 _TORCH_DTYPE_OP_NAMES = frozenset({"sum", "mean", "prod"})
64 _MS_POSITIONAL_DTYPE_OP_NAMES = frozenset({"SumExt", "MeanExt"})
66 def __init__(self, op_name, partial_type=None):
67 super().__init__(op_name)
68 if partial_type is None:
69 partial_type = ["sum"]
70 self.partial_type = partial_type
71 self._allow_partial_inputs = True
73 def preprocess(self, args: tuple, kwargs: dict) -> tuple:
74 """
75 Preprocess arguments for reduce operators.
77 Normalizes ``(input, dim, keepdim)`` across torch and MindSpore call
78 sites, converts DTensor inputs to local tensors, and builds
79 ``cache_values`` for layout inference.
81 Args:
82 args (tuple): Positional arguments passed to the operator call.
83 kwargs (dict): Keyword arguments passed to the operator call.
85 Returns:
86 tuple: ``(local_args, local_kwargs, cache_values)`` where
87 ``local_args`` is ``(input_local, dim, keepdim)`` and
88 ``cache_values`` is ``[input_layout, dim, keepdim]``.
89 """
90 has_dtype_arg = len(args) >= 4 or "dtype" in kwargs
91 normalized_args, _ = _normalize_reduce_args(*args, **kwargs)
92 input_tensor, dim, keepdim, dtype = normalized_args
93 local_args = (input_tensor.to_local(), dim, keepdim)
94 local_kwargs = {}
95 if has_dtype_arg:
96 if self.op_name in self._TORCH_DTYPE_OP_NAMES:
97 local_kwargs["dtype"] = dtype
98 elif self.op_name in self._MS_POSITIONAL_DTYPE_OP_NAMES:
99 local_args += (dtype,)
100 else:
101 raise TypeError(
102 f"For {self.op_name}, the `dtype` argument is not supported."
103 )
104 cache_values = [input_tensor.layout, dim, keepdim]
105 return local_args, local_kwargs, cache_values
107 def infer_layout(self, cache_values: list) -> Tuple[tuple, None]:
108 """
109 Infer output layout for reduce operators.
111 Rules:
112 1. Input must have a valid mesh_shape.
113 2. ``dim`` must be ``None``, ``int``, ``tuple[int]``, or ``list[int]`` —
114 never a ``Tensor``.
115 3. ``dim`` values must be valid axis indices for the input tensor.
116 4. Reduce axes are replaced with ``"None"`` (keepdim) or dropped;
117 Partial status is applied on sharded reduced axes.
119 Args:
120 cache_values (list): ``[input_layout, dim, keepdim]``.
122 Returns:
123 tuple: ``((output_layout,), None)``
125 Raises:
126 ValueError: If the input layout is missing or ``mesh_shape`` is ``None``.
127 TypeError: If ``dim`` is a ``Tensor``.
128 ValueError: If a ``dim`` axis index is out of range.
129 """
130 x_layout = cache_values[0]
131 dim = cache_values[1]
132 keepdim = cache_values[2]
134 if x_layout is None or x_layout.mesh_shape is None:
135 raise ValueError(
136 f"For {self.op_name}, input layout cannot be None."
137 )
139 # Check partial inputs
140 if not self._allow_partial_inputs:
141 self._check_partial_inputs([x_layout])
143 if dim is not None and not isinstance(dim, (int, tuple, list)):
144 raise TypeError(
145 f"For {self.op_name}, the `dim` argument should be `None`, `int`, "
146 f"`tuple[int]`, or `list[int]`, but got {type(dim).__name__}."
147 )
149 # Infer the output shape based on dim and keepdim
150 output_layout = self._infer_output_layout(x_layout, dim, keepdim)
152 return ((output_layout,), None)
154 def _infer_output_layout(self, x_layout, dim, keepdim):
155 """Infer output layout for reduce operator."""
156 # Case 1: dim is None — reduce all dimensions (scalar output or all-None when keepdim).
157 if dim is None:
158 return self._handle_all_axis_reduce(x_layout, keepdim)
160 # Case 2: dim is an empty tuple/list — reduce no dimensions, output layout equals input.
161 if isinstance(dim, (tuple, list)) and len(dim) == 0:
162 output_layout = Layout(
163 mesh_shape=x_layout.mesh_shape,
164 alias_name=x_layout.alias_name,
165 rank_list=x_layout.rank_list
166 )
167 return output_layout(*x_layout.alias_tensor_map)
169 # Case 3: dim is int, tuple, or list with at least one element.
170 output_layout = Layout(
171 mesh_shape=x_layout.mesh_shape,
172 alias_name=x_layout.alias_name,
173 rank_list=x_layout.rank_list
174 )
175 x_map = x_layout.alias_tensor_map
176 reduce_alias, x_map = self.replace_axis_with_none(dim, x_layout, keepdim)
177 output_layout = output_layout(*x_map)
178 self._apply_partial(output_layout, reduce_alias)
179 return output_layout
181 def _handle_all_axis_reduce(self, x_layout, keepdim):
182 """Handle the case where dim is None, meaning reduce all dimensions."""
183 layout = Layout(
184 mesh_shape=x_layout.mesh_shape,
185 alias_name=x_layout.alias_name,
186 rank_list=x_layout.rank_list
187 )
189 if not keepdim:
190 output_layout = layout()
191 else:
192 tensor_map = tuple(["None"] * len(x_layout.alias_tensor_map))
193 output_layout = layout(*tensor_map)
195 self._apply_partial(output_layout, x_layout.alias_tensor_map)
196 return output_layout
198 def replace_axis_with_none(self, dim, x_layout, keepdim):
199 """Replace or drop dimensions depending on keepdim."""
200 if not isinstance(dim, (tuple, list)):
201 dim = [dim]
202 else:
203 dim = list(dim)
205 rank = len(x_layout.alias_tensor_map)
206 for i, axis_id in enumerate(dim):
207 if axis_id < 0:
208 dim[i] = rank + axis_id
209 if not isinstance(axis_id, int) or dim[i] >= rank or dim[i] < 0:
210 raise ValueError(f"Invalid reduce axis index {axis_id} at position {i}.")
212 alias_tensor_map = x_layout.alias_tensor_map
213 reduce_alias = [alias_tensor_map[axis_id] for axis_id in dim if
214 alias_tensor_map[axis_id] is not None and alias_tensor_map[axis_id] != "None"]
215 reduce_alias = self._flatten_aliases(reduce_alias)
217 if keepdim:
218 return self._replace_keepdim(alias_tensor_map, reduce_alias)
219 return self._replace_dropdim(alias_tensor_map, reduce_alias, dim)
221 def _flatten_aliases(self, reduce_alias):
222 """Flatten reduce_alias into a list of atomic alias strings."""
223 flat = []
224 for alias in reduce_alias:
225 if isinstance(alias, (tuple, list)):
226 flat.extend(alias)
227 else:
228 flat.append(alias)
229 return flat
231 def _replace_keepdim(self, alias_tensor_map, reduce_alias):
232 """keepdim, replace reduce alias with 'None'."""
233 new_alias_map = []
234 for alias in alias_tensor_map:
235 if isinstance(alias, (tuple, list)):
236 new_alias = tuple("None" if item in reduce_alias else item for item in alias)
237 new_alias_map.append(new_alias)
238 else:
239 if alias in reduce_alias:
240 new_alias_map.append("None")
241 else:
242 new_alias_map.append(alias)
243 new_alias_map = self._compact_tensor_map(new_alias_map)
244 return reduce_alias, tuple(new_alias_map)
246 def _replace_dropdim(self, alias_tensor_map, reduce_alias, dim):
247 """Compress reduce dim."""
248 new_alias_map = []
249 for i, alias in enumerate(alias_tensor_map):
250 if i in dim:
251 continue
252 if isinstance(alias, (tuple, list)):
253 new_alias = tuple(item for item in alias if item not in reduce_alias)
254 if new_alias:
255 new_alias_map.append(new_alias)
256 else:
257 if alias in reduce_alias:
258 continue
259 new_alias_map.append(alias)
260 new_alias_map = self._compact_tensor_map(new_alias_map)
261 return reduce_alias, tuple(new_alias_map)
263 def _compact_tensor_map(self, alias_map: Sequence[StrOrTuple]) -> Tuple[StrOrTuple, ...]:
264 """Extend tensor map of 'None'."""
266 def _compress(elem: StrOrTuple) -> StrOrTuple:
267 if isinstance(elem, (list, tuple)):
268 compressed = tuple(_compress(e) for e in elem)
269 if len(compressed) == 1:
270 return compressed[0]
271 if all(x == 'None' for x in compressed):
272 return 'None'
273 return compressed
274 return elem
276 return tuple(_compress(elem) for elem in alias_map)
278 def _apply_partial(self, out_layout, alias):
279 """Apply all partial to given alias (string, tuple, list)."""
280 if alias == "None":
281 return
282 if isinstance(alias, (tuple, list)):
283 for elem in alias:
284 self._apply_partial(out_layout, elem)
285 else:
286 for ops in self.partial_type:
287 out_layout.set_partial_by_dev_axis(alias, ops)
290class SumExtDistributedOp(ReduceExtDistributedOpBase):
291 """Distributed implementation for SumExt operator."""
293 def __init__(self, op_name="SumExt"):
294 super().__init__(op_name, partial_type=["sum"])
297class MeanExtDistributedOp(ReduceExtDistributedOpBase):
298 """Distributed implementation for MeanExt operator."""
300 def __init__(self, op_name="MeanExt"):
301 super().__init__(op_name, partial_type=["avg"])
304class ReduceMaxDistributedOp(ReduceExtDistributedOpBase):
305 """Distributed implementation for ReduceMax operator."""
307 def __init__(self, op_name="ReduceMax"):
308 super().__init__(op_name, partial_type=["max"])
311class ProdExtDistributedOp(ReduceExtDistributedOpBase):
312 """
313 Distributed implementation for ProdExt operator (product of all elements or along a dim).
314 Compatible with torch.prod arguments.
315 """
317 def __init__(self, op_name="prod"):
318 super().__init__(op_name, partial_type=["prod"])
321class AllExtDistributedOp(ReduceExtDistributedOpBase):
322 """
323 Distributed implementation for All operator
324 Returns the cumulative sum of elements of input in the dimension dim.
325 """
327 def __init__(self, op_name="all"):
328 super().__init__(op_name, partial_type=["all"])
331class MaxDistributedOp(ReduceExtDistributedOpBase):
332 """
333 Distributed implementation for Pytorch style Max operator.
335 Supports three Pytorch behaviors:
336 1. torch.max(input) -> Global reduction (returns single Tensor)
337 2. torch.max(input, dim, keepdim=False) -> Dimension reduction (returns (values, indices))
338 3. torch.max(input, other) -> Element-wise max (returns single Tensor)
339 """
341 def __init__(self, op_name="max"):
342 super().__init__(op_name, partial_type=["max"])
344 def preprocess(self, args: tuple, kwargs: dict) -> tuple:
345 """
346 Preprocess arguments for Max/Min operators.
348 Routes between element-wise mode (two DTensor inputs) and reduction
349 mode (input, dim, keepdim). All parameters are positional with empty
350 kwargs for MindSpore Primitive compatibility.
352 Args:
353 args (tuple): Positional arguments passed to the operator call.
354 kwargs (dict): Keyword arguments passed to the operator call.
356 Returns:
357 tuple: ``(local_args, local_kwargs, cache_values)``.
358 """
359 # Element-wise mode: torch.max(a, b) — second arg is a DTensor.
360 if len(args) >= 2 and hasattr(args[1], "_layout"):
361 input_tensor = args[0]
362 other_tensor = args[1]
363 local_args = (input_tensor.to_local(), other_tensor.to_local())
364 local_kwargs = {}
365 cache_values = [input_tensor.layout, other_tensor.layout]
366 return local_args, local_kwargs, cache_values
368 # Reduction mode: torch.max(input, dim, keepdim) / MaxDim(input, dim, keepdim).
369 has_dtype_arg = len(args) >= 4 or "dtype" in kwargs
370 args, kwargs = _normalize_reduce_args(*args, **kwargs)
371 input_tensor, dim, keepdim, _ = args
372 if has_dtype_arg:
373 raise TypeError(
374 f"For {self.op_name}, the `dtype` argument is not supported."
375 )
376 # torch.max(x) global reduction: only pass the tensor so the call
377 # remains torch.max(local_x), not torch.max(local_x, None, False).
378 if dim is None:
379 local_args = (input_tensor.to_local(),)
380 else:
381 local_args = (input_tensor.to_local(), dim, keepdim)
382 local_kwargs = {}
384 cache_values = [input_tensor.layout, dim, keepdim]
385 return local_args, local_kwargs, cache_values
387 def infer_layout(self, cache_values: list) -> Tuple[tuple, None]:
388 """
389 Infer output layouts for Max/Min operators.
391 Rules:
392 1. Element-wise mode (two Layouts in *cache_values*): output layout
393 equals the first input layout.
394 2. Reduction mode: same rules as ``ReduceExtDistributedOpBase``.
395 3. When ``dim`` is ``None`` (global reduction), returns a single layout.
396 4. When ``dim`` is specified, returns ``(values_layout, indices_layout)``.
397 5. ``dim`` must be ``None``, ``int``, ``tuple[int]``, or ``list[int]``.
399 Args:
400 cache_values (list): ``[input_layout, dim, keepdim]`` for reduction,
401 or ``[input_layout, other_layout]`` for element-wise.
403 Returns:
404 tuple: ``((output_layout,), None)`` or ``((values_layout, indices_layout), None)``.
406 Raises:
407 ValueError: If the input layout is missing or ``mesh_shape`` is ``None``.
408 TypeError: If ``dim`` is not ``None``, ``int``, ``tuple``, or ``list``.
409 """
410 # Element-wise mode: two Layout objects in cache_values.
411 if len(cache_values) == 2 and hasattr(cache_values[1], "mesh_shape"):
412 # Check partial inputs
413 if not self._allow_partial_inputs:
414 self._check_partial_inputs(cache_values)
415 return ((deepcopy(cache_values[0]),), None)
417 x_layout = cache_values[0]
418 dim = cache_values[1]
419 keepdim = cache_values[2]
421 if x_layout is None or x_layout.mesh_shape is None:
422 raise ValueError(
423 f"For {self.op_name}, input layout cannot be None."
424 )
426 # Check partial inputs
427 if not self._allow_partial_inputs:
428 self._check_partial_inputs([x_layout])
430 if dim is not None and not isinstance(dim, (int, tuple, list)):
431 raise TypeError(
432 f"For {self.op_name}, the `dim` argument should be `None`, `int`, "
433 f"`tuple[int]`, or `list[int]`, but got {type(dim).__name__}."
434 )
436 values_layout = self._infer_output_layout(x_layout, dim, keepdim)
438 if dim is None:
439 # torch.max(input) -> Single Tensor
440 return ((values_layout,), None)
442 # torch.max(input, dim) -> (values, indices)
443 indices_layout = deepcopy(values_layout)
444 return ((values_layout, indices_layout), None)
447class MinDistributedOp(MaxDistributedOp):
448 """
449 Distributed implementation for Pytorch style Min operator.
451 Supports three Pytorch behaviors:
452 1. torch.min(input) -> Global reduction (returns single Tensor)
453 2. torch.min(input, dim, keepdim=False) -> Dimension reduction (returns (values, indices))
454 3. torch.min(input, other) -> Element-wise min (returns single Tensor)
455 """
457 def __init__(self, op_name="min"):
458 # Call the parent class (MaxDistributedOp) initialization
459 super().__init__(op_name=op_name)
460 # Override the partial_type to use "min" instead of "max" for the underlying communication
461 self.partial_type = ["min"]