Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / shard / ops / parallel_setitem.py: 38%
198 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 __setitem__ operator (PyTorch only).
18__setitem__ (tensor[key] = value) is an in-place operation.
19The LHS key follows the same classification as __getitem__:
20BASIC (int/slice/None/Ellipsis), ADVANCED (list/LongTensor), BOOL_MASK (not supported).
21"""
22from typing import Any, Tuple
24from hyper_parallel.core.dtensor.dtensor import DTensor, _build_layout
25from hyper_parallel.core.dtensor.placement_types import Replicate
26from hyper_parallel.platform import get_platform
27from .parallel_getitem import (
28 _BASIC,
29 _BOOL_MASK,
30 _broadcast_shapes,
31 _key_cache_descriptor,
32 _descriptor_to_expanded_actions,
33 _unwrap_key_for_local,
34 GetItemDistributedOp,
35)
36from .parallel_ops import DistributedOp
38platform = get_platform()
39Tensor = platform.Tensor
42def _normalize___setitem___args(self_t, key, value):
43 """Normalize __setitem__ arguments to canonical positional form.
45 __setitem__ is always called as func(self, key, value) with empty kwargs.
47 Args:
48 self_t: The DTensor self argument.
49 key: The indexing key.
50 value: The value to assign.
52 Returns:
53 tuple: ((self_t, key, value), {})
54 """
55 return (self_t, key, value), {}
58def _compute_lhs_shape(global_shape, expanded_actions, kind, op_name="__setitem__"):
59 """Compute the shape of the LHS slice that __getitem__ would produce.
61 This is used to validate that the RHS value can broadcast to the LHS slice.
63 Args:
64 global_shape: Tuple of self tensor's global shape.
65 expanded_actions: Expanded key actions from _descriptor_to_expanded_actions.
66 kind: "basic" or "advanced".
67 op_name: Operator name for error messages.
69 Returns:
70 tuple: Shape of the LHS slice.
71 """
72 if kind == _BASIC:
73 return _compute_basic_lhs_shape(global_shape, expanded_actions, op_name=op_name)
74 return _compute_advanced_lhs_shape(global_shape, expanded_actions, op_name=op_name)
77def _compute_basic_lhs_shape(global_shape, expanded_actions, op_name="__setitem__"):
78 """Compute LHS shape for basic indexing.
80 Raises ValueError if step is not None or 1.
81 """
82 lhs_dims = []
83 for action in expanded_actions:
84 action_type = action[0]
85 if action_type == "newaxis":
86 lhs_dims.append(1)
87 elif action_type == "slice":
88 start, stop, step = action[1], action[2], action[3]
89 step = step if step is not None else 1
90 if step != 1:
91 raise ValueError(
92 f"For {op_name}, slice step must be None or 1, but got {step}."
93 )
94 input_dim = action[4]
95 dim_size = global_shape[input_dim]
96 # Compute slice length
97 start_val = start if start is not None else 0
98 if start_val < 0:
99 start_val += dim_size
100 stop_val = stop if stop is not None else dim_size
101 if stop_val < 0:
102 stop_val += dim_size
103 # Clamp to valid range
104 start_val = max(0, min(start_val, dim_size))
105 stop_val = max(0, min(stop_val, dim_size))
106 slice_len = max(0, (stop_val - start_val + step - 1) // step)
107 lhs_dims.append(slice_len)
108 # int actions are skipped (dimension removed)
109 return tuple(lhs_dims) if lhs_dims else ()
112def _compute_slice_len(action, global_shape, op_name="__setitem__"):
113 """Compute output dim length for a slice action.
115 Raises ValueError if step is not None or 1.
116 """
117 start, stop, step = action[1], action[2], action[3]
118 step = step if step is not None else 1
119 if step != 1:
120 raise ValueError(
121 f"For {op_name}, slice step must be None or 1, but got {step}."
122 )
123 input_dim = action[4]
124 dim_size = global_shape[input_dim]
125 start_val = start if start is not None else 0
126 if start_val < 0:
127 start_val += dim_size
128 stop_val = stop if stop is not None else dim_size
129 if stop_val < 0:
130 stop_val += dim_size
131 start_val = max(0, min(start_val, dim_size))
132 stop_val = max(0, min(stop_val, dim_size))
133 return max(0, (stop_val - start_val + step - 1) // step)
136def _append_non_advanced_shape(action, global_shape, result, op_name="__setitem__"):
137 """Append output dim length for a non-advanced action (newaxis, slice; skip int)."""
138 if action[0] == "newaxis":
139 result.append(1)
140 elif action[0] == "slice":
141 result.append(_compute_slice_len(action, global_shape, op_name=op_name))
142 # int: dimension removed, skip
145def _compute_advanced_lhs_shape(global_shape, expanded_actions, op_name="__setitem__"):
146 """Compute LHS shape for advanced indexing.
148 Walks expanded_actions positionally: newaxis → 1, slice → slice_len,
149 int → skip, advanced block → replaced by broadcast shape B.
150 """
151 # Collect advanced index shapes and their positions
152 advanced_shapes = []
153 advanced_positions = []
154 for i, action in enumerate(expanded_actions):
155 if action[0] == "idx_list":
156 advanced_shapes.append((len(action[1]),))
157 advanced_positions.append(i)
158 elif action[0] == "idx_tensor":
159 advanced_shapes.append(tuple(action[1]["shape"]))
160 advanced_positions.append(i)
162 bcast_shape = _broadcast_shapes(advanced_shapes, op_name=op_name) if advanced_shapes else ()
164 advanced_pos_set = set(advanced_positions)
165 are_consecutive = (
166 len(advanced_positions) > 0
167 and advanced_positions == list(range(advanced_positions[0],
168 advanced_positions[-1] + 1))
169 )
171 result = []
173 if are_consecutive and advanced_positions:
174 first_adv_pos = advanced_positions[0]
175 last_adv_pos = advanced_positions[-1]
177 # Actions before the advanced block
178 for i in range(first_adv_pos):
179 _append_non_advanced_shape(expanded_actions[i], global_shape, result, op_name=op_name)
181 # B dims replace the advanced block
182 result.extend(bcast_shape)
184 # Actions after the advanced block
185 for i in range(last_adv_pos + 1, len(expanded_actions)):
186 _append_non_advanced_shape(expanded_actions[i], global_shape, result, op_name=op_name)
187 else:
188 # Non-consecutive: B goes at position 0
189 result.extend(bcast_shape)
191 for i, action in enumerate(expanded_actions):
192 if i in advanced_pos_set:
193 continue
194 _append_non_advanced_shape(action, global_shape, result)
196 return tuple(result) if result else ()
199def _compute_lhs_layout(self_layout, expanded_actions, kind, op_name="__setitem__"):
200 """Compute the LHS layout (getitem output layout) for the given key.
202 Args:
203 self_layout: Layout of self tensor.
204 expanded_actions: Expanded key actions.
205 kind: "basic" or "advanced".
206 op_name: Operator name for error messages.
208 Returns:
209 Layout: The layout that __getitem__ would produce for this key.
210 """
211 if kind == _BASIC:
212 return GetItemDistributedOp._infer_basic_output_layout( # pylint: disable=W0212
213 self_layout, expanded_actions
214 )
215 return GetItemDistributedOp._infer_advanced_output_layout( # pylint: disable=W0212
216 self_layout, expanded_actions, op_name=op_name
217 )
220def _build_broadcast_value_layout(mesh, lhs_layout, value_shape, lhs_shape):
221 """Build a layout for the value tensor with broadcast dims forced to Replicate.
223 When a value tensor has a smaller shape that needs to broadcast to the LHS
224 slice shape, any dimension where value size is 1 and LHS size > 1 must be
225 Replicate (cannot be sharded). Non-broadcast dims inherit LHS placements.
227 Args:
228 mesh: Device mesh.
229 lhs_layout: Expected LHS layout (getitem output layout).
230 value_shape: Shape of the value tensor.
231 lhs_shape: Expected LHS slice shape.
233 Returns:
234 Layout: Value layout adjusted for broadcasting.
235 """
236 ndim_diff = len(lhs_shape) - len(value_shape)
237 padded_value_shape = (1,) * ndim_diff + value_shape
239 placements = []
240 for v_size, l_size, lhs_placement in zip(
241 padded_value_shape, lhs_shape, lhs_layout.alias_placements
242 ):
243 if v_size == 1 and l_size > 1:
244 placements.append(Replicate())
245 else:
246 placements.append(lhs_placement)
248 return _build_layout(mesh, tuple(placements), len(placements))
251class SetItemDistributedOp(DistributedOp):
252 """Distributed implementation for tensor.__setitem__.
254 __setitem__ (tensor[key] = value) is an in-place operation.
255 The LHS key follows the same validation rules as __getitem__.
256 RHS value can be a scalar, Tensor, or DTensor.
258 Sharding constraints:
259 - Same LHS constraints as __getitem__ (sharded dims cannot be indexed).
260 - BoolTensor masks are rejected.
261 - RHS DTensor must have mesh consistent with self and broadcast-compatible
262 layout (non-broadcast dims must match LHS layout, broadcast dims must be
263 Replicate).
264 - RHS Tensor/scalar must be broadcastable to LHS slice shape.
265 """
267 def preprocess(self, args: tuple, kwargs: dict) -> tuple:
268 """Preprocess arguments for __setitem__.
270 Args:
271 args: (self_tensor, key, value)
272 kwargs: Empty dict for __setitem__.
274 Returns:
275 tuple: (local_args, local_kwargs, cache_values)
276 """
277 norm_args = _normalize___setitem___args(*args, **kwargs)[0]
278 self_t, key, value = norm_args
280 local_key, key_cache, lhs_shape, lhs_layout = self._preprocess_setitem_key(
281 self_t, key, self.op_name
282 )
284 local_value, value_desc = self._preprocess_setitem_value(
285 value, key_cache[3], lhs_shape, lhs_layout, key_cache[0].mesh, self.op_name,
286 )
288 cache_values = key_cache + [value_desc, lhs_shape]
289 return ((self_t.to_local(), local_key, local_value), {}, cache_values)
291 @staticmethod
292 def _preprocess_setitem_key(self_t, key, op_name):
293 """Process key and compute LHS metadata for __setitem__ preprocess.
295 Args:
296 self_t: The self DTensor.
297 key: Raw indexing key.
298 op_name: Operator name for error messages.
300 Returns:
301 tuple: (local_key, key_cache, lhs_shape, lhs_layout)
302 key_cache is [self_layout, key_desc, global_shape, kind].
303 """
304 self_layout = self_t.layout
305 global_shape = tuple(self_t.shape)
306 key_desc, kind = _key_cache_descriptor(key, op_name=op_name)
307 local_key = _unwrap_key_for_local(key)
308 key_cache = [self_layout, key_desc, global_shape, kind]
310 if kind == _BOOL_MASK:
311 return local_key, key_cache, None, None
313 expanded_actions = _descriptor_to_expanded_actions(
314 key_desc, len(global_shape), op_name=op_name
315 )
316 lhs_shape = _compute_lhs_shape(global_shape, expanded_actions, kind, op_name=op_name)
317 lhs_layout = _compute_lhs_layout(self_layout, expanded_actions, kind, op_name=op_name)
319 return local_key, key_cache, lhs_shape, lhs_layout
321 @staticmethod
322 def _preprocess_setitem_value(value, kind, lhs_shape, lhs_layout, self_mesh, op_name):
323 """Process the value argument for __setitem__ preprocess.
325 Handles DTensor unwrapping, plain-tensor broadcast validation and sharding,
326 and scalar passthrough.
328 Args:
329 value: The raw value tensor (DTensor, Tensor, or scalar).
330 kind: "basic", "advanced", or "bool_mask".
331 lhs_shape: Computed LHS slice shape (tuple).
332 lhs_layout: Computed LHS slice Layout.
333 self_mesh: DeviceMesh of the self tensor.
334 op_name: Operator name for error messages.
336 Returns:
337 (local_value, value_desc): Processed local tensor and cache descriptor.
338 """
339 if isinstance(value, DTensor):
340 value_shape = tuple(value.shape)
341 local_value = value.to_local()
342 if not value_shape:
343 return local_value, None # 0-D DTensor treated as scalar
344 return local_value, ("dtensor", value.layout, value_shape)
346 if isinstance(value, Tensor):
347 value_desc = ("plain_tensor", tuple(value.shape))
348 if value.ndim > 0 and kind != _BOOL_MASK:
349 _validate_value_broadcast(op_name, tuple(value.shape), lhs_shape)
350 # pylint: disable=C0415
351 from hyper_parallel.core.dtensor.layout import _get_slice_tensor_by_layout
352 value_layout = _build_broadcast_value_layout(
353 self_mesh, lhs_layout, tuple(value.shape), lhs_shape
354 )
355 ndim_diff = len(lhs_shape) - value.ndim
356 if ndim_diff > 0:
357 padded_value = value.reshape((1,) * ndim_diff + tuple(value.shape))
358 else:
359 padded_value = value
360 local_value = _get_slice_tensor_by_layout(padded_value, value_layout)
361 return local_value, value_desc
362 return value, value_desc
364 return value, None
366 def infer_layout(self, cache_values: list) -> Tuple[tuple, None]: # pylint: disable=W0221
367 """Infer output layout for __setitem__.
369 __setitem__ is in-place, so output layout equals self layout.
370 Validates LHS key constraints and RHS value compatibility.
372 Rules:
373 1. self must not be in Partial status.
374 2. BoolTensor mask indexing is not supported.
375 3. Any dimension being written must be replicated.
376 4. slice step must be None or 1.
377 5. Advanced index tensors must be replicated.
378 6. RHS DTensor must have mesh consistent with self and
379 broadcast-compatible layout (non-broadcast dims match LHS
380 layout, broadcast dims must be Replicate).
381 7. RHS Tensor/scalar must be broadcastable to LHS slice shape.
382 8. Output layout equals self_layout (in-place).
384 Args:
385 cache_values: [self_layout, key_desc, global_shape, kind, value_desc, lhs_shape]
387 Returns:
388 tuple: ((self_layout,), None)
390 Raises:
391 ValueError: If any constraint is violated.
392 """
393 self_layout = cache_values[0]
394 key_desc = cache_values[1]
395 global_shape = cache_values[2]
396 kind = cache_values[3]
397 value_desc = cache_values[4]
398 lhs_shape = cache_values[5]
400 # key_desc -> expanded_actions (for LHS validation)
401 expanded_actions = _descriptor_to_expanded_actions(
402 key_desc, len(global_shape), op_name=self.op_name
403 )
405 if not self._allow_partial_inputs:
406 self._check_partial_inputs([self_layout])
407 self._validate_input_layouts(
408 self_layout, expanded_actions, global_shape, kind, value_desc, lhs_shape
409 )
410 return ((self_layout,), None)
412 @staticmethod
413 def _validate_input_layouts(self_layout, expanded_actions, global_shape, kind, value_desc, lhs_shape):
414 """Validate sharding constraints for __setitem__."""
415 op_name = "__setitem__"
417 GetItemDistributedOp._validate_input_layouts( # pylint: disable=W0212
418 self_layout, expanded_actions, global_shape, kind, op_name=op_name
419 )
421 if value_desc is None:
422 return
424 val_kind = value_desc[0]
425 if val_kind == "plain_tensor":
426 _validate_value_broadcast(op_name, value_desc[1], lhs_shape)
427 elif val_kind == "dtensor":
428 SetItemDistributedOp._validate_dtensor_value(
429 self_layout, expanded_actions, kind, value_desc[1], value_desc[2], lhs_shape, op_name,
430 )
431 else:
432 raise ValueError(
433 f"For {op_name}, unsupported value descriptor kind: {val_kind}."
434 )
436 @staticmethod
437 def _validate_dtensor_value(
438 self_layout, expanded_actions, kind, val_layout, val_shape, lhs_shape, op_name):
439 """Validate DTensor RHS mesh and layout."""
440 _validate_value_broadcast(op_name, val_shape, lhs_shape)
442 if val_layout.mesh.to_hash() != self_layout.mesh.to_hash():
443 raise ValueError(
444 f"For {op_name}, value DTensor must be on same mesh as self, "
445 f"but got mesh {val_layout.mesh_shape} "
446 f"vs {self_layout.mesh_shape}."
447 )
449 expected_layout = _compute_lhs_layout(
450 self_layout, expanded_actions, kind, op_name=op_name
451 )
452 SetItemDistributedOp._validate_value_layout(
453 val_layout, expected_layout, val_shape, lhs_shape, op_name,
454 )
456 @staticmethod
457 def _validate_value_layout(val_layout, expected_layout, val_shape, lhs_shape, op_name):
458 """Validate RHS DTensor placement against expected LHS placement."""
459 ndim_diff = len(lhs_shape) - len(val_shape)
460 val_alias = val_layout.alias_tensor_map
461 expected_alias = expected_layout.alias_tensor_map
463 for i, v_size in enumerate(val_shape):
464 lhs_index = ndim_diff + i
465 e_size = lhs_shape[lhs_index]
466 v_alias = val_alias[i]
467 e_alias = expected_alias[lhs_index]
469 if v_size == 1 and e_size > 1:
470 if v_alias != "None":
471 raise ValueError(
472 f"For {op_name}, value broadcast dim {i} (size {v_size}) "
473 f"must be Replicate, but got {v_alias}."
474 )
475 elif v_alias != e_alias:
476 raise ValueError(
477 f"For {op_name}, value layout mismatch at dim {i}: "
478 f"expected {e_alias}, but got {v_alias}."
479 )
481 def wrap_output(self, py_output: Any, output_layouts: tuple) -> None:
482 """Override wrap_output for __setitem__.
484 __setitem__ is in-place and returns None in PyTorch.
485 The default wrap_output would try to wrap None into a DTensor and fail.
487 Args:
488 py_output: The output from the local function call (None).
489 output_layouts: Output layouts from infer_layout.
491 Returns:
492 None: __setitem__ has no return value.
493 """
494 return None
497def _validate_value_broadcast(op_name, val_shape, lhs_shape):
498 """Validate that value shape can broadcast to LHS slice shape.
500 Args:
501 op_name: Operator name for error messages.
502 val_shape: Shape of RHS value.
503 lhs_shape: Expected LHS slice shape.
505 Raises:
506 ValueError: If shapes are not broadcastable.
507 """
508 if len(val_shape) > len(lhs_shape):
509 raise ValueError(
510 f"For {op_name}, value shape {val_shape} cannot broadcast "
511 f"to target shape {lhs_shape}."
512 )
514 # Pad val_shape with leading 1s
515 padded = (1,) * (len(lhs_shape) - len(val_shape)) + val_shape
516 for v, t in zip(padded, lhs_shape):
517 if v not in (1, t):
518 raise ValueError(
519 f"For {op_name}, value shape {val_shape} cannot broadcast "
520 f"to target shape {lhs_shape}."
521 )