Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / shard / ops / parallel_gather.py: 71%
240 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-08-04 05:18 +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 Gather operator.
17"""
19from typing import Tuple
21from hyper_parallel.core.dtensor.layout import Layout
22from hyper_parallel.platform import get_platform
23from .parallel_ops import DistributedOp
26def _normalize_index_select_args(input_tensor, dim, index):
27 return (input_tensor, dim, index), {}
30def _normalize_gatherd_args(input_tensor, dim, index, **kwargs):
31 """Normalize torch.gather / GatherD args, forwarding keyword-only params."""
32 return (input_tensor, dim, index), kwargs
35def _normalize_gathernd_args(input_tensor, indices):
36 return (input_tensor, indices), {}
39class IndexSelectDistributedOp(DistributedOp):
40 """Distributed implementation for Index Select operator."""
42 def preprocess(self, args: tuple, kwargs: dict) -> tuple:
43 """
44 Preprocess arguments for IndexSelect operator.
46 Args:
47 args (tuple): Input arguments (input, dim, index).
48 kwargs (dict): Keyword arguments.
50 Returns:
51 tuple: (local_args, local_kwargs, cache_values)
52 """
53 args, _ = _normalize_index_select_args(*args, **kwargs)
54 input_tensor, dim, index = args[0], args[1], args[2]
55 local_args = (input_tensor.to_local(), dim, index.to_local())
56 local_kwargs = {}
57 cache_values = [input_tensor.layout, index.layout, dim]
58 return local_args, local_kwargs, cache_values
60 def infer_layout(self, cache_values: list) -> Tuple[tuple, None]: # pylint: disable=W0221
61 """
62 Infer output layouts for Index Select operations.
64 Rules:
65 1. Input and index must not have Partial status.
66 2. cache_values must contain [input_layout, index_layout, dim].
67 3. dim must be within the valid input rank range.
68 4. index must be one-dimensional.
69 5. Output replaces the selected input dimension with the index layout.
70 6. If the selected input dimension is sharded, output carries Partial('sum').
72 Args:
73 cache_values (list): [input_layout, index_layout, dim].
75 Returns:
76 tuple: ((output_layout,), None)
78 Raises:
79 ValueError: If input layouts are not compatible or have partial status.
80 """
81 if not self._allow_partial_inputs:
82 self._check_partial_inputs(cache_values[:2])
84 # Parse layout info
85 p_layout, i_layout, axis = cache_values[0], cache_values[1], cache_values[2]
87 p_tensor_map = p_layout.alias_tensor_map
88 i_tensor_map = i_layout.alias_tensor_map
90 # 1. Validate the axis range before any manipulation
91 if axis < -len(p_tensor_map) or axis >= len(p_tensor_map):
92 raise ValueError(
93 f"For {self.op_name}, dim value {axis} is out of valid range"
94 )
96 # 2. Convert negative axis to positive index to avoid Python slicing bugs
97 if axis < 0:
98 axis += len(p_tensor_map)
100 if len(i_tensor_map) != 1:
101 raise ValueError(
102 f"For {self.op_name}, index is not a one-dimensional Tensor"
103 )
105 # 3. Create output layout map
106 # We allow sharding on the `axis`. Since `index_select` replaces the `axis`
107 # dimension with the `index` dimension, if `axis` was sharded, that mesh
108 # dimension is removed from the output tensor map.
109 output_tensor_map = list(p_tensor_map[:axis]) + list(i_tensor_map) + list(p_tensor_map[axis + 1 :])
111 output_layout = Layout(
112 mesh_shape=p_layout.mesh_shape,
113 alias_name=p_layout.alias_name,
114 rank_list=p_layout.rank_list,
115 )
116 output_layout = output_layout(*output_tensor_map)
118 # 4. Implicit Communication via Partial Layout
119 # If the gather axis was sharded, the local output will only be a masked partial result.
120 # We set the output layout to Partial('sum') for that specific mesh dimension so the
121 # OpDispatcher handles the AllReduce automatically when this tensor is used later.
122 shard_mesh_dim_name = p_tensor_map[axis]
123 if shard_mesh_dim_name != "None":
124 # Handle possible multi-axis sharding tuple
125 if isinstance(shard_mesh_dim_name, tuple):
126 for dim_name in shard_mesh_dim_name:
127 if dim_name != "None":
128 output_layout.set_partial_by_dev_axis(dim_name, 'sum')
129 else:
130 output_layout.set_partial_by_dev_axis(shard_mesh_dim_name, 'sum')
132 return ((output_layout,), None)
134 def get_expand_impl(self, func, infer_result, cache_values): # pylint: disable=W0221
135 """
136 Get the expanded execution implementation for Index Select.
137 """
138 p_layout = cache_values[0]
139 axis = cache_values[2]
140 if axis < 0:
141 axis += len(p_layout.alias_tensor_map)
143 shard_mesh_dim_name = p_layout.alias_tensor_map[axis]
145 # If the axis is NOT sharded, fallback to standard execution
146 if shard_mesh_dim_name == "None":
147 return None
149 # If the axis IS sharded, return a custom function with Masking ONLY.
150 # The explicit AllReduce is completely removed.
151 def expand_impl(input_tensor, dim, index, **kwargs):
152 platform = get_platform()
153 mesh = p_layout.mesh
155 # Fetch the communication group for the sharded mesh dimension
156 if isinstance(shard_mesh_dim_name, tuple):
157 target_dim_name = next(d for d in shard_mesh_dim_name if d != "None")
158 else:
159 target_dim_name = shard_mesh_dim_name
161 comm_group_info = mesh.get_comm_group_by_axis(target_dim_name)
162 group = comm_group_info.group if hasattr(comm_group_info, 'group') else comm_group_info
164 # Get the rank of the current device within this specific communication group
165 group_rank = platform.get_group_local_rank(group=group)
167 # Calculate global index boundaries for the local chunk
168 local_dim_size = input_tensor.shape[dim]
169 start_idx = group_rank * local_dim_size
170 end_idx = start_idx + local_dim_size
172 # 1. Compute mask: True for indices that belong to the current rank
173 mask = (index >= start_idx) & (index < end_idx)
175 # 2. Shift global indices to local indices
176 safe_index = index - start_idx
178 # Clamp safe_index to valid local ranges to prevent CUDA out-of-bounds
179 # errors during the local index_select (invalid ones will be masked out anyway).
180 safe_index = safe_index.clamp(min=0, max=local_dim_size - 1)
182 # 3. Perform local index_select using tensor's built-in method
183 local_out = input_tensor.index_select(dim, safe_index, **kwargs)
185 # 4. Mask out the invalid indices (set them to 0)
186 # Reshape the 1D mask to broadcast against the output shape
187 mask_shape = [1] * local_out.ndim
188 mask_shape[dim] = -1
189 mask_reshaped = mask.reshape(mask_shape).to(local_out.dtype)
191 local_out = local_out * mask_reshaped
193 # Return the partial local tensor directly. The framework's layout engine
194 # and OpDispatcher will trigger the AllReduce when this Partial tensor
195 # is redistributed to a non-partial layout.
196 return local_out
198 return expand_impl
201class GatherDDistributedOp(DistributedOp):
202 """Distributed implementation for GatherD operator.
204 GatherD gathers values along a specified axis from the input tensor using the index tensor.
206 Signature: GatherD(input, dim, index) -> output
208 Key constraints:
209 - Input and index must have the same number of dimensions
210 - Output inherits the sharding pattern of the input tensor
211 """
213 def preprocess(self, args: tuple, kwargs: dict) -> tuple:
214 """
215 Preprocess arguments for GatherD operator.
217 Args:
218 args (tuple): Input arguments (input, dim, index).
219 kwargs (dict): Keyword arguments.
221 Returns:
222 tuple: (local_args, local_kwargs, cache_values)
223 """
224 args, kwargs = _normalize_gatherd_args(*args, **kwargs)
225 input_tensor, dim, index = args[0], args[1], args[2]
226 local_args = (input_tensor.to_local(), dim, index.to_local())
227 local_kwargs = kwargs
228 cache_values = [input_tensor.layout, index.layout, dim]
229 return local_args, local_kwargs, cache_values
231 def infer_layout(self, cache_values: list) -> Tuple[tuple, None]: # pylint: disable=W0221
232 """
233 Infer output layouts for GatherD operations.
235 Rules:
236 1. Input and index must not have Partial status.
237 2. cache_values must contain [input_layout, index_layout, dim].
238 3. Input and index must have the same rank.
239 4. dim must be within the valid input rank range.
240 5. Input and index must use the same sharding on non-dim axes.
241 6. Output inherits index layout and becomes Partial('sum') when dim is sharded.
243 Args:
244 cache_values (list): [input_layout, index_layout, dim].
246 Returns:
247 tuple: ((output_layout,), None)
249 Raises:
250 ValueError: If input layouts are not compatible or have partial status.
251 """
252 if not self._allow_partial_inputs:
253 self._check_partial_inputs(cache_values[:2])
255 input_layout, index_layout, dim = cache_values[0], cache_values[1], cache_values[2]
256 # Validate layouts exist
257 if input_layout is None or not hasattr(input_layout, "tensor_map"):
258 raise ValueError(f"For {self.op_name}, input layout cannot be None")
259 if index_layout is None or not hasattr(index_layout, "tensor_map"):
260 raise ValueError(f"For {self.op_name}, index layout cannot be None")
261 input_tensor_map = input_layout.alias_tensor_map
262 index_tensor_map = index_layout.alias_tensor_map
263 # Validate same rank
264 if len(input_tensor_map) != len(index_tensor_map):
265 raise ValueError(
266 f"For {self.op_name}, input and index must have the same number of dimensions. "
267 f"Got input rank={len(input_tensor_map)}, index rank={len(index_tensor_map)}"
268 )
269 # Validate dim is in valid range
270 rank = len(input_tensor_map)
271 if dim < -rank or dim >= rank:
272 raise ValueError(
273 f"For {self.op_name}, dim value {dim} is out of valid range [{-rank}, {rank-1}]"
274 )
275 # Normalize negative dim
276 if dim < 0:
277 dim = dim + rank
278 for axis, (input_axis_map, index_axis_map) in enumerate(zip(input_tensor_map, index_tensor_map)):
279 if axis == dim:
280 continue
281 if input_axis_map != index_axis_map:
282 raise ValueError(
283 f"For {self.op_name}, input and index must use the same sharding on non-dim axis {axis}. "
284 f"Got input tensor_map={input_tensor_map}, index tensor_map={index_tensor_map}, dim={dim}"
285 )
286 # Output inherits index layout
287 output_layout = Layout(
288 mesh_shape=index_layout.mesh_shape,
289 alias_name=index_layout.alias_name,
290 rank_list=index_layout.rank_list,
291 )
292 output_layout.set_tensor_map(index_layout.tensor_map)
293 dim_axis_name = input_tensor_map[dim]
294 if dim_axis_name != "None":
295 # pylint: disable=protected-access
296 # Inherit current partial state from index layout
297 output_layout._partial = list(index_layout.partial)
298 if isinstance(dim_axis_name, tuple):
299 for axis_name in dim_axis_name:
300 if axis_name != "None":
301 output_layout.set_partial_by_dev_axis(axis_name, 'sum')
302 else:
303 output_layout.set_partial_by_dev_axis(dim_axis_name, 'sum')
304 # pylint: disable=protected-access
305 # Rebuild readable alias tensor map
306 output_layout._alias_tensor_map = output_layout._build_readable_tensor_map()
307 # pylint: disable=protected-access
308 # Sync tensor_map to placement representation
309 output_layout.tensor_map_to_placement()
310 # Update compact string description
311 output_layout.update_compact_str()
312 return ((output_layout,), None)
314 def get_expand_impl(self, func, infer_result, cache_values): # pylint: disable=W0221
315 """
316 Returns the execution implementation wrapper for distributed GatherD.
318 When the dim axis is sharded, each rank gathers from its local slice of the input tensor.
319 The indices need to be adjusted to account for the local partition offset.
321 Args:
322 func: The original GatherD function to wrap
323 infer_result: The inferred output layouts and extra info
324 cache_values: [input_layout, index_layout, dim]
326 Returns:
327 Callable: Distributed implementation wrapper, or None if no sharding
328 """
329 input_layout = cache_values[0]
330 dim = cache_values[2]
331 if dim < 0:
332 dim += len(input_layout.tensor_map)
333 input_alias_map = input_layout.alias_tensor_map
334 # Check if dim axis is sharded (enhanced MP)
335 if input_alias_map[dim] == "None": # native sharding, no need for custom implementation
336 return None
338 dim_axis_name = input_alias_map[dim]
339 if isinstance(dim_axis_name, tuple):
340 dim_axis_name = next(axis for axis in dim_axis_name if axis != "None")
342 def distributed_gatherd_impl(*args, **kwargs):
343 """
344 Distributed GatherD implementation for sharded dim axis.
346 Each rank gathers from its local slice of input tensor.
347 Indices are adjusted by subtracting the local partition offset.
348 """
349 input_tensor = args[0]
350 index_tensor = args[2]
351 # Calculate local partition offset for the dim axis
352 mesh = input_layout.mesh
353 mesh_dim_idx = input_layout.alias_name.index(dim_axis_name)
354 # Get the coordinate of current rank along the mesh dimension
355 dim_coord = mesh.get_local_rank(mesh_dim_idx)
356 # Calculate the size of input tensor's dim dimension per partition
357 input_dim_size = input_tensor.shape[dim]
358 # Calculate the starting index of local partition
359 local_start_index = int(dim_coord * input_dim_size)
360 local_end_index = int(local_start_index + input_dim_size)
361 # Adjust indices: subtract local_start_index to map global indices to local range
362 # This is similar to how Embedding shifts indices for Row Parallelism
363 adjusted_index = index_tensor - local_start_index
364 # Create mask to identify out-of-bounds indices
365 # Indices outside [0, local_dim_size) belong to other partitions
366 mask = (index_tensor >= local_start_index) & (index_tensor < local_end_index)
367 # Cross-platform cast to matching int dtype
368 mask_int = mask.to(index_tensor.dtype)
369 # Zero out invalid indices to prevent out-of-bounds access
370 safe_index = adjusted_index * mask_int
371 # Replace original index tensor with adjusted index
372 new_args = list(args)
373 new_args[2] = safe_index
374 # Execute native GatherD with adjusted indices
375 output = func(*new_args, **kwargs)
376 # Zero out outputs corresponding to invalid indices
377 mask_int = mask_int.to(output.dtype)
378 output = output * mask_int
379 return output
380 return distributed_gatherd_impl
383class GatherNdDistributedOp(DistributedOp):
384 """Distributed implementation for GatherNd operator."""
386 def preprocess(self, args: tuple, kwargs: dict) -> tuple:
387 """
388 Preprocess arguments for GatherNd operator.
390 NOTE: aclop packed-args normalization (for MindSpore aclop operators
391 that pack args as ``(prim, name, (real_args...))``) is handled
392 upstream in ``OpDispatcher._dispatch_layout_infer`` via
393 ``_normalize_aclop_args``. This method receives clean unpacked args.
395 Args:
396 args (tuple): Input arguments (input, indices).
397 kwargs (dict): Keyword arguments.
399 Returns:
400 tuple: (local_args, local_kwargs, cache_values)
401 """
402 args, _ = _normalize_gathernd_args(*args, **kwargs)
403 input_tensor, indices = args[0], args[1]
404 local_input = input_tensor.to_local() if hasattr(input_tensor, "_layout") else input_tensor
405 local_indices = indices.to_local()
406 local_args = (local_input, local_indices)
407 local_kwargs = {}
408 cache_values = [
409 input_tensor.layout if hasattr(input_tensor, "_layout") else None,
410 indices.layout,
411 input_tensor.shape,
412 indices.shape,
413 ]
414 return local_args, local_kwargs, cache_values
416 def infer_layout(self, cache_values: list) -> Tuple[tuple, None]: # pylint: disable=W0221
417 """
418 Infer output layout for GatherNd.
420 Rules:
421 1. Input and indices must not have Partial status.
422 2. cache_values must contain [input_layout_or_None, indices_layout, input_shape, indices_shape].
423 3. indices[-1] (K) must be replicated.
424 4. input indexed dims [0:K) must be replicated when input layout is provided.
425 5. Output inherits indices[:-1] sharding plus input trailing dims input[K:].
427 For GatherNd: out.shape = indices.shape[:-1] + input_x.shape[K:], where K = indices.shape[-1].
429 This implementation:
430 - Inherits sharding from indices[:-1].
431 - Allows sharding on input_x trailing dims input_x[K:].
432 - Requires input_x[:K] to be replicated ("None") if input_layout is provided.
433 - Requires indices[-1] (K dim) to be replicated ("None").
435 Output Layout:
436 output_tensor_map = indices_tensor_map[:-1] + input_tensor_map[K:]
437 If input_layout is None, input trailing dims are treated as replicated ("None").
439 Args:
440 cache_values (list): [input_layout_or_None, indices_layout, input_shape, indices_shape].
442 Returns:
443 tuple: ((output_layout,), None)
445 Raises:
446 ValueError: If input layouts, tensor maps, or shapes violate the rules above.
447 """
448 input_layout, indices_layout = self._parse_input_layouts(cache_values[:2])
449 if not self._allow_partial_inputs:
450 self._check_partial_inputs([input_layout, indices_layout])
452 input_shape, indices_shape = self._get_input_shapes(cache_values[2:])
453 k, trail_rank = self._get_k_and_trailing_rank(input_shape, indices_shape)
455 input_tensor_map, indices_tensor_map = self._validate_tensor_maps(
456 input_layout, indices_layout, k
457 )
459 # Output sharding: inherit indices[:-1] + input_x[K:].
460 if input_tensor_map is None:
461 output_tensor_map = tuple(indices_tensor_map[:-1]) + ("None",) * trail_rank
462 else:
463 output_tensor_map = tuple(indices_tensor_map[:-1]) + tuple(input_tensor_map[k:])
465 output_layout = Layout(
466 mesh_shape=indices_layout.mesh_shape,
467 alias_name=indices_layout.alias_name,
468 rank_list=indices_layout.rank_list,
469 )
471 if output_tensor_map:
472 output_layout = output_layout(*output_tensor_map)
473 else:
474 output_layout = output_layout("None")
476 return ((output_layout,), None)
478 def _parse_input_layouts(self, layouts):
479 """Parse and validate input layouts."""
480 if len(layouts) < 2:
481 raise ValueError(
482 f"For {self.op_name}, requires at least 2 input layouts, but got {len(layouts)}"
483 )
485 input_layout, indices_layout = layouts[0], layouts[1]
487 # Extra inputs are allowed only when they are non-tensor args (layout is None).
488 for extra_layout in layouts[2:]:
489 if extra_layout is not None:
490 raise ValueError(
491 f"For {self.op_name}, only supports 2 tensor inputs, but got extra tensor layout: "
492 f"{extra_layout}"
493 )
495 # For GatherNd: input_layout can be None (treated as fully replicated), but indices_layout must exist.
496 if indices_layout is None or not hasattr(indices_layout, "alias_tensor_map"):
497 raise ValueError(f"For {self.op_name}, indices layout cannot be None")
499 return input_layout, indices_layout
501 def _validate_tensor_maps(self, input_layout, indices_layout, k):
502 """Validate tensor maps constraints for GatherNd."""
503 indices_tensor_map = indices_layout.alias_tensor_map
505 # Validate: indices tensor_map must exist and last dimension cannot be split.
506 if not indices_tensor_map:
507 raise ValueError(f"For {self.op_name}, indices tensor_map cannot be empty")
509 last_axis = indices_tensor_map[-1]
510 if not self._is_none_axis(last_axis):
511 raise ValueError(
512 f"For {self.op_name}, the last dimension of indices cannot be split. "
513 f"Got indices[-1] = {last_axis}"
514 )
516 # Validate input only when layout is provided.
517 input_tensor_map = None
518 if input_layout is not None:
519 input_tensor_map = input_layout.alias_tensor_map
521 if k > len(input_tensor_map):
522 raise ValueError(
523 f"For {self.op_name}, indices last dim (K={k}) is larger than input rank "
524 f"({len(input_tensor_map)})"
525 )
527 # Indexed dims [0:K) must be replicated.
528 for axis_name in input_tensor_map[:k]:
529 if not self._is_none_axis(axis_name):
530 raise ValueError(
531 f"For {self.op_name}, input_x cannot be split on indexed dims [0:{k}). "
532 f"These dims must be 'None', but got tensor_map: {input_tensor_map}"
533 )
535 return input_tensor_map, indices_tensor_map
537 def _get_input_shapes(self, shape_values):
538 """Get input and indices shapes from cache values."""
539 input_shapes = None
540 if shape_values and len(shape_values) == 2:
541 input_shapes = shape_values
543 if input_shapes is None:
544 raise ValueError(
545 f"For {self.op_name}, missing input_shapes in cache_values."
546 )
548 input_shape = input_shapes[0]
549 indices_shape = input_shapes[1]
550 if input_shape is None or indices_shape is None:
551 raise ValueError(f"For {self.op_name}, input_shapes contains None: {input_shapes}")
553 input_shape = self._normalize_shape(input_shape, "input")
554 indices_shape = self._normalize_shape(indices_shape, "indices")
556 if len(indices_shape) < 1:
557 raise ValueError(f"For {self.op_name}, indices shape invalid: {indices_shape}")
559 return input_shape, indices_shape
561 def _normalize_shape(self, shape, name):
562 """Normalize shape-like object to tuple of int."""
563 try:
564 norm = tuple(shape)
565 except TypeError as err:
566 raise ValueError(f"For {self.op_name}, {name} shape is not iterable: {shape}") from err
568 try:
569 norm = tuple(int(dim) for dim in norm)
570 except (TypeError, ValueError) as err:
571 raise ValueError(f"For {self.op_name}, {name} shape contains non-integer dims: {norm}") from err
573 return norm
575 def _get_k_and_trailing_rank(self, input_shape, indices_shape):
576 """Compute K and trailing rank = len(input_shape) - K, where K is indices_shape[-1]."""
577 k = indices_shape[-1]
578 try:
579 k = int(k)
580 except (TypeError, ValueError) as err:
581 raise ValueError(f"For {self.op_name}, indices last dim (K) is invalid: {k}") from err
583 if k <= 0:
584 raise ValueError(f"For {self.op_name}, indices last dim (K) must be positive, but got {k}")
586 trail_rank = len(input_shape) - k
587 if trail_rank < 0:
588 raise ValueError(
589 f"For {self.op_name}, indices last dim (K={k}) is larger than input rank ({len(input_shape)})"
590 )
592 return k, trail_rank
594 def _is_none_axis(self, axis_name):
595 """
596 Check if an axis name represents no sharding.
597 """
598 if axis_name == "None":
599 return True
601 if isinstance(axis_name, tuple):
602 return all(name == "None" for name in axis_name)
604 return False