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-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 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):
31 return (input_tensor, dim, index), {}
34def _normalize_gathernd_args(input_tensor, indices):
35 return (input_tensor, indices), {}
38class IndexSelectDistributedOp(DistributedOp):
39 """Distributed implementation for Index Select operator."""
41 def preprocess(self, args: tuple, kwargs: dict) -> tuple:
42 """
43 Preprocess arguments for IndexSelect operator.
45 Args:
46 args (tuple): Input arguments (input, dim, index).
47 kwargs (dict): Keyword arguments.
49 Returns:
50 tuple: (local_args, local_kwargs, cache_values)
51 """
52 args, _ = _normalize_index_select_args(*args, **kwargs)
53 input_tensor, dim, index = args[0], args[1], args[2]
54 local_args = (input_tensor.to_local(), dim, index.to_local())
55 local_kwargs = {}
56 cache_values = [input_tensor.layout, index.layout, dim]
57 return local_args, local_kwargs, cache_values
59 def infer_layout(self, cache_values: list) -> Tuple[tuple, None]: # pylint: disable=W0221
60 """
61 Infer output layouts for Index Select operations.
63 Rules:
64 1. Input and index must not have Partial status.
65 2. cache_values must contain [input_layout, index_layout, dim].
66 3. dim must be within the valid input rank range.
67 4. index must be one-dimensional.
68 5. Output replaces the selected input dimension with the index layout.
69 6. If the selected input dimension is sharded, output carries Partial('sum').
71 Args:
72 cache_values (list): [input_layout, index_layout, dim].
74 Returns:
75 tuple: ((output_layout,), None)
77 Raises:
78 ValueError: If input layouts are not compatible or have partial status.
79 """
80 if not self._allow_partial_inputs:
81 self._check_partial_inputs(cache_values[:2])
83 # Parse layout info
84 p_layout, i_layout, axis = cache_values[0], cache_values[1], cache_values[2]
86 p_tensor_map = p_layout.alias_tensor_map
87 i_tensor_map = i_layout.alias_tensor_map
89 # 1. Validate the axis range before any manipulation
90 if axis < -len(p_tensor_map) or axis >= len(p_tensor_map):
91 raise ValueError(
92 f"For {self.op_name}, dim value {axis} is out of valid range"
93 )
95 # 2. Convert negative axis to positive index to avoid Python slicing bugs
96 if axis < 0:
97 axis += len(p_tensor_map)
99 if len(i_tensor_map) != 1:
100 raise ValueError(
101 f"For {self.op_name}, index is not a one-dimensional Tensor"
102 )
104 # 3. Create output layout map
105 # We allow sharding on the `axis`. Since `index_select` replaces the `axis`
106 # dimension with the `index` dimension, if `axis` was sharded, that mesh
107 # dimension is removed from the output tensor map.
108 output_tensor_map = list(p_tensor_map[:axis]) + list(i_tensor_map) + list(p_tensor_map[axis + 1 :])
110 output_layout = Layout(
111 mesh_shape=p_layout.mesh_shape,
112 alias_name=p_layout.alias_name,
113 rank_list=p_layout.rank_list,
114 )
115 output_layout = output_layout(*output_tensor_map)
117 # 4. Implicit Communication via Partial Layout
118 # If the gather axis was sharded, the local output will only be a masked partial result.
119 # We set the output layout to Partial('sum') for that specific mesh dimension so the
120 # OpDispatcher handles the AllReduce automatically when this tensor is used later.
121 shard_mesh_dim_name = p_tensor_map[axis]
122 if shard_mesh_dim_name != "None":
123 # Handle possible multi-axis sharding tuple
124 if isinstance(shard_mesh_dim_name, tuple):
125 for dim_name in shard_mesh_dim_name:
126 if dim_name != "None":
127 output_layout.set_partial_by_dev_axis(dim_name, 'sum')
128 else:
129 output_layout.set_partial_by_dev_axis(shard_mesh_dim_name, 'sum')
131 return ((output_layout,), None)
133 def get_expand_impl(self, func, infer_result, cache_values): # pylint: disable=W0221
134 """
135 Get the expanded execution implementation for Index Select.
136 """
137 p_layout = cache_values[0]
138 axis = cache_values[2]
139 if axis < 0:
140 axis += len(p_layout.alias_tensor_map)
142 shard_mesh_dim_name = p_layout.alias_tensor_map[axis]
144 # If the axis is NOT sharded, fallback to standard execution
145 if shard_mesh_dim_name == "None":
146 return None
148 # If the axis IS sharded, return a custom function with Masking ONLY.
149 # The explicit AllReduce is completely removed.
150 def expand_impl(input_tensor, dim, index, **kwargs):
151 platform = get_platform()
152 mesh = p_layout.mesh
154 # Fetch the communication group for the sharded mesh dimension
155 if isinstance(shard_mesh_dim_name, tuple):
156 target_dim_name = next(d for d in shard_mesh_dim_name if d != "None")
157 else:
158 target_dim_name = shard_mesh_dim_name
160 comm_group_info = mesh.get_comm_group_by_axis(target_dim_name)
161 group = comm_group_info.group if hasattr(comm_group_info, 'group') else comm_group_info
163 # Get the rank of the current device within this specific communication group
164 group_rank = platform.get_group_local_rank(group=group)
166 # Calculate global index boundaries for the local chunk
167 local_dim_size = input_tensor.shape[dim]
168 start_idx = group_rank * local_dim_size
169 end_idx = start_idx + local_dim_size
171 # 1. Compute mask: True for indices that belong to the current rank
172 mask = (index >= start_idx) & (index < end_idx)
174 # 2. Shift global indices to local indices
175 safe_index = index - start_idx
177 # Clamp safe_index to valid local ranges to prevent CUDA out-of-bounds
178 # errors during the local index_select (invalid ones will be masked out anyway).
179 safe_index = safe_index.clamp(min=0, max=local_dim_size - 1)
181 # 3. Perform local index_select using tensor's built-in method
182 local_out = input_tensor.index_select(dim, safe_index, **kwargs)
184 # 4. Mask out the invalid indices (set them to 0)
185 # Reshape the 1D mask to broadcast against the output shape
186 mask_shape = [1] * local_out.ndim
187 mask_shape[dim] = -1
188 mask_reshaped = mask.reshape(mask_shape).to(local_out.dtype)
190 local_out = local_out * mask_reshaped
192 # Return the partial local tensor directly. The framework's layout engine
193 # and OpDispatcher will trigger the AllReduce when this Partial tensor
194 # is redistributed to a non-partial layout.
195 return local_out
197 return expand_impl
200class GatherDDistributedOp(DistributedOp):
201 """Distributed implementation for GatherD operator.
203 GatherD gathers values along a specified axis from the input tensor using the index tensor.
205 Signature: GatherD(input, dim, index) -> output
207 Key constraints:
208 - Input and index must have the same number of dimensions
209 - Output inherits the sharding pattern of the input tensor
210 """
212 def preprocess(self, args: tuple, kwargs: dict) -> tuple:
213 """
214 Preprocess arguments for GatherD operator.
216 Args:
217 args (tuple): Input arguments (input, dim, index).
218 kwargs (dict): Keyword arguments.
220 Returns:
221 tuple: (local_args, local_kwargs, cache_values)
222 """
223 args, _ = _normalize_gatherd_args(*args, **kwargs)
224 input_tensor, dim, index = args[0], args[1], args[2]
225 local_args = (input_tensor.to_local(), dim, index.to_local())
226 local_kwargs = {}
227 cache_values = [input_tensor.layout, index.layout, dim]
228 return local_args, local_kwargs, cache_values
230 def infer_layout(self, cache_values: list) -> Tuple[tuple, None]: # pylint: disable=W0221
231 """
232 Infer output layouts for GatherD operations.
234 Rules:
235 1. Input and index must not have Partial status.
236 2. cache_values must contain [input_layout, index_layout, dim].
237 3. Input and index must have the same rank.
238 4. dim must be within the valid input rank range.
239 5. Input and index must use the same sharding on non-dim axes.
240 6. Output inherits index layout and becomes Partial('sum') when dim is sharded.
242 Args:
243 cache_values (list): [input_layout, index_layout, dim].
245 Returns:
246 tuple: ((output_layout,), None)
248 Raises:
249 ValueError: If input layouts are not compatible or have partial status.
250 """
251 if not self._allow_partial_inputs:
252 self._check_partial_inputs(cache_values[:2])
254 input_layout, index_layout, dim = cache_values[0], cache_values[1], cache_values[2]
255 # Validate layouts exist
256 if input_layout is None or not hasattr(input_layout, "tensor_map"):
257 raise ValueError(f"For {self.op_name}, input layout cannot be None")
258 if index_layout is None or not hasattr(index_layout, "tensor_map"):
259 raise ValueError(f"For {self.op_name}, index layout cannot be None")
260 input_tensor_map = input_layout.alias_tensor_map
261 index_tensor_map = index_layout.alias_tensor_map
262 # Validate same rank
263 if len(input_tensor_map) != len(index_tensor_map):
264 raise ValueError(
265 f"For {self.op_name}, input and index must have the same number of dimensions. "
266 f"Got input rank={len(input_tensor_map)}, index rank={len(index_tensor_map)}"
267 )
268 # Validate dim is in valid range
269 rank = len(input_tensor_map)
270 if dim < -rank or dim >= rank:
271 raise ValueError(
272 f"For {self.op_name}, dim value {dim} is out of valid range [{-rank}, {rank-1}]"
273 )
274 # Normalize negative dim
275 if dim < 0:
276 dim = dim + rank
277 for axis, (input_axis_map, index_axis_map) in enumerate(zip(input_tensor_map, index_tensor_map)):
278 if axis == dim:
279 continue
280 if input_axis_map != index_axis_map:
281 raise ValueError(
282 f"For {self.op_name}, input and index must use the same sharding on non-dim axis {axis}. "
283 f"Got input tensor_map={input_tensor_map}, index tensor_map={index_tensor_map}, dim={dim}"
284 )
285 # Output inherits index layout
286 output_layout = Layout(
287 mesh_shape=index_layout.mesh_shape,
288 alias_name=index_layout.alias_name,
289 rank_list=index_layout.rank_list,
290 )
291 output_layout.set_tensor_map(index_layout.tensor_map)
292 dim_axis_name = input_tensor_map[dim]
293 if dim_axis_name != "None":
294 # pylint: disable=protected-access
295 # Inherit current partial state from index layout
296 output_layout._partial = list(index_layout.partial)
297 if isinstance(dim_axis_name, tuple):
298 for axis_name in dim_axis_name:
299 if axis_name != "None":
300 output_layout.set_partial_by_dev_axis(axis_name, 'sum')
301 else:
302 output_layout.set_partial_by_dev_axis(dim_axis_name, 'sum')
303 # pylint: disable=protected-access
304 # Rebuild readable alias tensor map
305 output_layout._alias_tensor_map = output_layout._build_readable_tensor_map()
306 # pylint: disable=protected-access
307 # Sync tensor_map to placement representation
308 output_layout.tensor_map_to_placement()
309 # Update compact string description
310 output_layout.update_compact_str()
311 return ((output_layout,), None)
313 def get_expand_impl(self, func, infer_result, cache_values): # pylint: disable=W0221
314 """
315 Returns the execution implementation wrapper for distributed GatherD.
317 When the dim axis is sharded, each rank gathers from its local slice of the input tensor.
318 The indices need to be adjusted to account for the local partition offset.
320 Args:
321 func: The original GatherD function to wrap
322 infer_result: The inferred output layouts and extra info
323 cache_values: [input_layout, index_layout, dim]
325 Returns:
326 Callable: Distributed implementation wrapper, or None if no sharding
327 """
328 input_layout = cache_values[0]
329 dim = cache_values[2]
330 if dim < 0:
331 dim += len(input_layout.tensor_map)
332 input_alias_map = input_layout.alias_tensor_map
333 # Check if dim axis is sharded (enhanced MP)
334 if input_alias_map[dim] == "None": # native sharding, no need for custom implementation
335 return None
337 dim_axis_name = input_alias_map[dim]
338 if isinstance(dim_axis_name, tuple):
339 dim_axis_name = next(axis for axis in dim_axis_name if axis != "None")
341 def distributed_gatherd_impl(*args, **kwargs):
342 """
343 Distributed GatherD implementation for sharded dim axis.
345 Each rank gathers from its local slice of input tensor.
346 Indices are adjusted by subtracting the local partition offset.
347 """
348 input_tensor = args[0]
349 index_tensor = args[2]
350 # Calculate local partition offset for the dim axis
351 mesh = input_layout.mesh
352 mesh_dim_idx = input_layout.alias_name.index(dim_axis_name)
353 # Get the coordinate of current rank along the mesh dimension
354 dim_coord = mesh.get_local_rank(mesh_dim_idx)
355 # Calculate the size of input tensor's dim dimension per partition
356 input_dim_size = input_tensor.shape[dim]
357 # Calculate the starting index of local partition
358 local_start_index = int(dim_coord * input_dim_size)
359 local_end_index = int(local_start_index + input_dim_size)
360 # Adjust indices: subtract local_start_index to map global indices to local range
361 # This is similar to how Embedding shifts indices for Row Parallelism
362 adjusted_index = index_tensor - local_start_index
363 # Create mask to identify out-of-bounds indices
364 # Indices outside [0, local_dim_size) belong to other partitions
365 mask = (index_tensor >= local_start_index) & (index_tensor < local_end_index)
366 # Cross-platform cast to matching int dtype
367 mask_int = mask.to(index_tensor.dtype)
368 # Zero out invalid indices to prevent out-of-bounds access
369 safe_index = adjusted_index * mask_int
370 # Replace original index tensor with adjusted index
371 new_args = list(args)
372 new_args[2] = safe_index
373 # Execute native GatherD with adjusted indices
374 output = func(*new_args, **kwargs)
375 # Zero out outputs corresponding to invalid indices
376 mask_int = mask_int.to(output.dtype)
377 output = output * mask_int
378 return output
379 return distributed_gatherd_impl
382class GatherNdDistributedOp(DistributedOp):
383 """Distributed implementation for GatherNd operator."""
385 def preprocess(self, args: tuple, kwargs: dict) -> tuple:
386 """
387 Preprocess arguments for GatherNd operator.
389 NOTE: aclop packed-args normalization (for MindSpore aclop operators
390 that pack args as ``(prim, name, (real_args...))``) is handled
391 upstream in ``OpDispatcher._dispatch_layout_infer`` via
392 ``_normalize_aclop_args``. This method receives clean unpacked args.
394 Args:
395 args (tuple): Input arguments (input, indices).
396 kwargs (dict): Keyword arguments.
398 Returns:
399 tuple: (local_args, local_kwargs, cache_values)
400 """
401 args, _ = _normalize_gathernd_args(*args, **kwargs)
402 input_tensor, indices = args[0], args[1]
403 local_input = input_tensor.to_local() if hasattr(input_tensor, "_layout") else input_tensor
404 local_indices = indices.to_local()
405 local_args = (local_input, local_indices)
406 local_kwargs = {}
407 cache_values = [
408 input_tensor.layout if hasattr(input_tensor, "_layout") else None,
409 indices.layout,
410 input_tensor.shape,
411 indices.shape,
412 ]
413 return local_args, local_kwargs, cache_values
415 def infer_layout(self, cache_values: list) -> Tuple[tuple, None]: # pylint: disable=W0221
416 """
417 Infer output layout for GatherNd.
419 Rules:
420 1. Input and indices must not have Partial status.
421 2. cache_values must contain [input_layout_or_None, indices_layout, input_shape, indices_shape].
422 3. indices[-1] (K) must be replicated.
423 4. input indexed dims [0:K) must be replicated when input layout is provided.
424 5. Output inherits indices[:-1] sharding plus input trailing dims input[K:].
426 For GatherNd: out.shape = indices.shape[:-1] + input_x.shape[K:], where K = indices.shape[-1].
428 This implementation:
429 - Inherits sharding from indices[:-1].
430 - Allows sharding on input_x trailing dims input_x[K:].
431 - Requires input_x[:K] to be replicated ("None") if input_layout is provided.
432 - Requires indices[-1] (K dim) to be replicated ("None").
434 Output Layout:
435 output_tensor_map = indices_tensor_map[:-1] + input_tensor_map[K:]
436 If input_layout is None, input trailing dims are treated as replicated ("None").
438 Args:
439 cache_values (list): [input_layout_or_None, indices_layout, input_shape, indices_shape].
441 Returns:
442 tuple: ((output_layout,), None)
444 Raises:
445 ValueError: If input layouts, tensor maps, or shapes violate the rules above.
446 """
447 input_layout, indices_layout = self._parse_input_layouts(cache_values[:2])
448 if not self._allow_partial_inputs:
449 self._check_partial_inputs([input_layout, indices_layout])
451 input_shape, indices_shape = self._get_input_shapes(cache_values[2:])
452 k, trail_rank = self._get_k_and_trailing_rank(input_shape, indices_shape)
454 input_tensor_map, indices_tensor_map = self._validate_tensor_maps(
455 input_layout, indices_layout, k
456 )
458 # Output sharding: inherit indices[:-1] + input_x[K:].
459 if input_tensor_map is None:
460 output_tensor_map = tuple(indices_tensor_map[:-1]) + ("None",) * trail_rank
461 else:
462 output_tensor_map = tuple(indices_tensor_map[:-1]) + tuple(input_tensor_map[k:])
464 output_layout = Layout(
465 mesh_shape=indices_layout.mesh_shape,
466 alias_name=indices_layout.alias_name,
467 rank_list=indices_layout.rank_list,
468 )
470 if output_tensor_map:
471 output_layout = output_layout(*output_tensor_map)
472 else:
473 output_layout = output_layout("None")
475 return ((output_layout,), None)
477 def _parse_input_layouts(self, layouts):
478 """Parse and validate input layouts."""
479 if len(layouts) < 2:
480 raise ValueError(
481 f"For {self.op_name}, requires at least 2 input layouts, but got {len(layouts)}"
482 )
484 input_layout, indices_layout = layouts[0], layouts[1]
486 # Extra inputs are allowed only when they are non-tensor args (layout is None).
487 for extra_layout in layouts[2:]:
488 if extra_layout is not None:
489 raise ValueError(
490 f"For {self.op_name}, only supports 2 tensor inputs, but got extra tensor layout: "
491 f"{extra_layout}"
492 )
494 # For GatherNd: input_layout can be None (treated as fully replicated), but indices_layout must exist.
495 if indices_layout is None or not hasattr(indices_layout, "alias_tensor_map"):
496 raise ValueError(f"For {self.op_name}, indices layout cannot be None")
498 return input_layout, indices_layout
500 def _validate_tensor_maps(self, input_layout, indices_layout, k):
501 """Validate tensor maps constraints for GatherNd."""
502 indices_tensor_map = indices_layout.alias_tensor_map
504 # Validate: indices tensor_map must exist and last dimension cannot be split.
505 if not indices_tensor_map:
506 raise ValueError(f"For {self.op_name}, indices tensor_map cannot be empty")
508 last_axis = indices_tensor_map[-1]
509 if not self._is_none_axis(last_axis):
510 raise ValueError(
511 f"For {self.op_name}, the last dimension of indices cannot be split. "
512 f"Got indices[-1] = {last_axis}"
513 )
515 # Validate input only when layout is provided.
516 input_tensor_map = None
517 if input_layout is not None:
518 input_tensor_map = input_layout.alias_tensor_map
520 if k > len(input_tensor_map):
521 raise ValueError(
522 f"For {self.op_name}, indices last dim (K={k}) is larger than input rank "
523 f"({len(input_tensor_map)})"
524 )
526 # Indexed dims [0:K) must be replicated.
527 for axis_name in input_tensor_map[:k]:
528 if not self._is_none_axis(axis_name):
529 raise ValueError(
530 f"For {self.op_name}, input_x cannot be split on indexed dims [0:{k}). "
531 f"These dims must be 'None', but got tensor_map: {input_tensor_map}"
532 )
534 return input_tensor_map, indices_tensor_map
536 def _get_input_shapes(self, shape_values):
537 """Get input and indices shapes from cache values."""
538 input_shapes = None
539 if shape_values and len(shape_values) == 2:
540 input_shapes = shape_values
542 if input_shapes is None:
543 raise ValueError(
544 f"For {self.op_name}, missing input_shapes in cache_values."
545 )
547 input_shape = input_shapes[0]
548 indices_shape = input_shapes[1]
549 if input_shape is None or indices_shape is None:
550 raise ValueError(f"For {self.op_name}, input_shapes contains None: {input_shapes}")
552 input_shape = self._normalize_shape(input_shape, "input")
553 indices_shape = self._normalize_shape(indices_shape, "indices")
555 if len(indices_shape) < 1:
556 raise ValueError(f"For {self.op_name}, indices shape invalid: {indices_shape}")
558 return input_shape, indices_shape
560 def _normalize_shape(self, shape, name):
561 """Normalize shape-like object to tuple of int."""
562 try:
563 norm = tuple(shape)
564 except TypeError as err:
565 raise ValueError(f"For {self.op_name}, {name} shape is not iterable: {shape}") from err
567 try:
568 norm = tuple(int(dim) for dim in norm)
569 except (TypeError, ValueError) as err:
570 raise ValueError(f"For {self.op_name}, {name} shape contains non-integer dims: {norm}") from err
572 return norm
574 def _get_k_and_trailing_rank(self, input_shape, indices_shape):
575 """Compute K and trailing rank = len(input_shape) - K, where K is indices_shape[-1]."""
576 k = indices_shape[-1]
577 try:
578 k = int(k)
579 except (TypeError, ValueError) as err:
580 raise ValueError(f"For {self.op_name}, indices last dim (K) is invalid: {k}") from err
582 if k <= 0:
583 raise ValueError(f"For {self.op_name}, indices last dim (K) must be positive, but got {k}")
585 trail_rank = len(input_shape) - k
586 if trail_rank < 0:
587 raise ValueError(
588 f"For {self.op_name}, indices last dim (K={k}) is larger than input rank ({len(input_shape)})"
589 )
591 return k, trail_rank
593 def _is_none_axis(self, axis_name):
594 """
595 Check if an axis name represents no sharding.
596 """
597 if axis_name == "None":
598 return True
600 if isinstance(axis_name, tuple):
601 return all(name == "None" for name in axis_name)
603 return False