Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / shard / ops / parallel_setitem.py: 37%
196 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +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
24import torch
26from hyper_parallel.core.dtensor.dtensor import DTensor, _build_layout
27from hyper_parallel.core.dtensor.placement_types import Replicate
28from .parallel_getitem import (
29 _BASIC,
30 _BOOL_MASK,
31 _broadcast_shapes,
32 _key_cache_descriptor,
33 _descriptor_to_expanded_actions,
34 _unwrap_key_for_local,
35 GetItemDistributedOp,
36)
37from .parallel_ops import DistributedOp
40def _normalize___setitem___args(self_t, key, value):
41 """Normalize __setitem__ arguments to canonical positional form.
43 __setitem__ is always called as func(self, key, value) with empty kwargs.
45 Args:
46 self_t: The DTensor self argument.
47 key: The indexing key.
48 value: The value to assign.
50 Returns:
51 tuple: ((self_t, key, value), {})
52 """
53 return (self_t, key, value), {}
56def _compute_lhs_shape(global_shape, expanded_actions, kind, op_name="__setitem__"):
57 """Compute the shape of the LHS slice that __getitem__ would produce.
59 This is used to validate that the RHS value can broadcast to the LHS slice.
61 Args:
62 global_shape: Tuple of self tensor's global shape.
63 expanded_actions: Expanded key actions from _descriptor_to_expanded_actions.
64 kind: "basic" or "advanced".
65 op_name: Operator name for error messages.
67 Returns:
68 tuple: Shape of the LHS slice.
69 """
70 if kind == _BASIC:
71 return _compute_basic_lhs_shape(global_shape, expanded_actions, op_name=op_name)
72 return _compute_advanced_lhs_shape(global_shape, expanded_actions, op_name=op_name)
75def _compute_basic_lhs_shape(global_shape, expanded_actions, op_name="__setitem__"):
76 """Compute LHS shape for basic indexing.
78 Raises ValueError if step is not None or 1.
79 """
80 lhs_dims = []
81 for action in expanded_actions:
82 action_type = action[0]
83 if action_type == "newaxis":
84 lhs_dims.append(1)
85 elif action_type == "slice":
86 start, stop, step = action[1], action[2], action[3]
87 step = step if step is not None else 1
88 if step != 1:
89 raise ValueError(
90 f"For {op_name}, slice step must be None or 1, but got {step}."
91 )
92 input_dim = action[4]
93 dim_size = global_shape[input_dim]
94 # Compute slice length
95 start_val = start if start is not None else 0
96 if start_val < 0:
97 start_val += dim_size
98 stop_val = stop if stop is not None else dim_size
99 if stop_val < 0:
100 stop_val += dim_size
101 # Clamp to valid range
102 start_val = max(0, min(start_val, dim_size))
103 stop_val = max(0, min(stop_val, dim_size))
104 slice_len = max(0, (stop_val - start_val + step - 1) // step)
105 lhs_dims.append(slice_len)
106 # int actions are skipped (dimension removed)
107 return tuple(lhs_dims) if lhs_dims else ()
110def _compute_slice_len(action, global_shape, op_name="__setitem__"):
111 """Compute output dim length for a slice action.
113 Raises ValueError if step is not None or 1.
114 """
115 start, stop, step = action[1], action[2], action[3]
116 step = step if step is not None else 1
117 if step != 1:
118 raise ValueError(
119 f"For {op_name}, slice step must be None or 1, but got {step}."
120 )
121 input_dim = action[4]
122 dim_size = global_shape[input_dim]
123 start_val = start if start is not None else 0
124 if start_val < 0:
125 start_val += dim_size
126 stop_val = stop if stop is not None else dim_size
127 if stop_val < 0:
128 stop_val += dim_size
129 start_val = max(0, min(start_val, dim_size))
130 stop_val = max(0, min(stop_val, dim_size))
131 return max(0, (stop_val - start_val + step - 1) // step)
134def _append_non_advanced_shape(action, global_shape, result, op_name="__setitem__"):
135 """Append output dim length for a non-advanced action (newaxis, slice; skip int)."""
136 if action[0] == "newaxis":
137 result.append(1)
138 elif action[0] == "slice":
139 result.append(_compute_slice_len(action, global_shape, op_name=op_name))
140 # int: dimension removed, skip
143def _compute_advanced_lhs_shape(global_shape, expanded_actions, op_name="__setitem__"):
144 """Compute LHS shape for advanced indexing.
146 Walks expanded_actions positionally: newaxis → 1, slice → slice_len,
147 int → skip, advanced block → replaced by broadcast shape B.
148 """
149 # Collect advanced index shapes and their positions
150 advanced_shapes = []
151 advanced_positions = []
152 for i, action in enumerate(expanded_actions):
153 if action[0] == "idx_list":
154 advanced_shapes.append((len(action[1]),))
155 advanced_positions.append(i)
156 elif action[0] == "idx_tensor":
157 advanced_shapes.append(tuple(action[1]["shape"]))
158 advanced_positions.append(i)
160 bcast_shape = _broadcast_shapes(advanced_shapes, op_name=op_name) if advanced_shapes else ()
162 advanced_pos_set = set(advanced_positions)
163 are_consecutive = (
164 len(advanced_positions) > 0
165 and advanced_positions == list(range(advanced_positions[0],
166 advanced_positions[-1] + 1))
167 )
169 result = []
171 if are_consecutive and advanced_positions:
172 first_adv_pos = advanced_positions[0]
173 last_adv_pos = advanced_positions[-1]
175 # Actions before the advanced block
176 for i in range(first_adv_pos):
177 _append_non_advanced_shape(expanded_actions[i], global_shape, result, op_name=op_name)
179 # B dims replace the advanced block
180 result.extend(bcast_shape)
182 # Actions after the advanced block
183 for i in range(last_adv_pos + 1, len(expanded_actions)):
184 _append_non_advanced_shape(expanded_actions[i], global_shape, result, op_name=op_name)
185 else:
186 # Non-consecutive: B goes at position 0
187 result.extend(bcast_shape)
189 for i, action in enumerate(expanded_actions):
190 if i in advanced_pos_set:
191 continue
192 _append_non_advanced_shape(action, global_shape, result)
194 return tuple(result) if result else ()
197def _compute_lhs_layout(self_layout, expanded_actions, kind, op_name="__setitem__"):
198 """Compute the LHS layout (getitem output layout) for the given key.
200 Args:
201 self_layout: Layout of self tensor.
202 expanded_actions: Expanded key actions.
203 kind: "basic" or "advanced".
204 op_name: Operator name for error messages.
206 Returns:
207 Layout: The layout that __getitem__ would produce for this key.
208 """
209 if kind == _BASIC:
210 return GetItemDistributedOp._infer_basic_output_layout( # pylint: disable=W0212
211 self_layout, expanded_actions
212 )
213 return GetItemDistributedOp._infer_advanced_output_layout( # pylint: disable=W0212
214 self_layout, expanded_actions, op_name=op_name
215 )
218def _build_broadcast_value_layout(mesh, lhs_layout, value_shape, lhs_shape):
219 """Build a layout for the value tensor with broadcast dims forced to Replicate.
221 When a value tensor has a smaller shape that needs to broadcast to the LHS
222 slice shape, any dimension where value size is 1 and LHS size > 1 must be
223 Replicate (cannot be sharded). Non-broadcast dims inherit LHS placements.
225 Args:
226 mesh: Device mesh.
227 lhs_layout: Expected LHS layout (getitem output layout).
228 value_shape: Shape of the value tensor.
229 lhs_shape: Expected LHS slice shape.
231 Returns:
232 Layout: Value layout adjusted for broadcasting.
233 """
234 ndim_diff = len(lhs_shape) - len(value_shape)
235 padded_value_shape = (1,) * ndim_diff + value_shape
237 placements = []
238 for v_size, l_size, lhs_placement in zip(
239 padded_value_shape, lhs_shape, lhs_layout.alias_placements
240 ):
241 if v_size == 1 and l_size > 1:
242 placements.append(Replicate())
243 else:
244 placements.append(lhs_placement)
246 return _build_layout(mesh, tuple(placements), len(placements))
249class SetItemDistributedOp(DistributedOp):
250 """Distributed implementation for tensor.__setitem__.
252 __setitem__ (tensor[key] = value) is an in-place operation.
253 The LHS key follows the same validation rules as __getitem__.
254 RHS value can be a scalar, Tensor, or DTensor.
256 Sharding constraints:
257 - Same LHS constraints as __getitem__ (sharded dims cannot be indexed).
258 - BoolTensor masks are rejected.
259 - RHS DTensor must have mesh consistent with self and broadcast-compatible
260 layout (non-broadcast dims must match LHS layout, broadcast dims must be
261 Replicate).
262 - RHS Tensor/scalar must be broadcastable to LHS slice shape.
263 """
265 def preprocess(self, args: tuple, kwargs: dict) -> tuple:
266 """Preprocess arguments for __setitem__.
268 Args:
269 args: (self_tensor, key, value)
270 kwargs: Empty dict for __setitem__.
272 Returns:
273 tuple: (local_args, local_kwargs, cache_values)
274 """
275 norm_args = _normalize___setitem___args(*args, **kwargs)[0]
276 self_t, key, value = norm_args
278 local_key, key_cache, lhs_shape, lhs_layout = self._preprocess_setitem_key(
279 self_t, key, self.op_name
280 )
282 local_value, value_desc = self._preprocess_setitem_value(
283 value, key_cache[3], lhs_shape, lhs_layout, key_cache[0].mesh, self.op_name,
284 )
286 cache_values = key_cache + [value_desc, lhs_shape]
287 return ((self_t.to_local(), local_key, local_value), {}, cache_values)
289 @staticmethod
290 def _preprocess_setitem_key(self_t, key, op_name):
291 """Process key and compute LHS metadata for __setitem__ preprocess.
293 Args:
294 self_t: The self DTensor.
295 key: Raw indexing key.
296 op_name: Operator name for error messages.
298 Returns:
299 tuple: (local_key, key_cache, lhs_shape, lhs_layout)
300 key_cache is [self_layout, key_desc, global_shape, kind].
301 """
302 self_layout = self_t.layout
303 global_shape = tuple(self_t.shape)
304 key_desc, kind = _key_cache_descriptor(key, op_name=op_name)
305 local_key = _unwrap_key_for_local(key)
306 key_cache = [self_layout, key_desc, global_shape, kind]
308 if kind == _BOOL_MASK:
309 return local_key, key_cache, None, None
311 expanded_actions = _descriptor_to_expanded_actions(
312 key_desc, len(global_shape), op_name=op_name
313 )
314 lhs_shape = _compute_lhs_shape(global_shape, expanded_actions, kind, op_name=op_name)
315 lhs_layout = _compute_lhs_layout(self_layout, expanded_actions, kind, op_name=op_name)
317 return local_key, key_cache, lhs_shape, lhs_layout
319 @staticmethod
320 def _preprocess_setitem_value(value, kind, lhs_shape, lhs_layout, self_mesh, op_name):
321 """Process the value argument for __setitem__ preprocess.
323 Handles DTensor unwrapping, plain-tensor broadcast validation and sharding,
324 and scalar passthrough.
326 Args:
327 value: The raw value tensor (DTensor, Tensor, or scalar).
328 kind: "basic", "advanced", or "bool_mask".
329 lhs_shape: Computed LHS slice shape (tuple).
330 lhs_layout: Computed LHS slice Layout.
331 self_mesh: DeviceMesh of the self tensor.
332 op_name: Operator name for error messages.
334 Returns:
335 (local_value, value_desc): Processed local tensor and cache descriptor.
336 """
337 if isinstance(value, DTensor):
338 value_shape = tuple(value.shape)
339 local_value = value.to_local()
340 if not value_shape:
341 return local_value, None # 0-D DTensor treated as scalar
342 return local_value, ("dtensor", value.layout, value_shape)
344 if isinstance(value, torch.Tensor):
345 value_desc = ("plain_tensor", tuple(value.shape))
346 if value.ndim > 0 and kind != _BOOL_MASK:
347 _validate_value_broadcast(op_name, tuple(value.shape), lhs_shape)
348 # pylint: disable=C0415
349 from hyper_parallel.core.dtensor.layout import _get_slice_tensor_by_layout
350 value_layout = _build_broadcast_value_layout(
351 self_mesh, lhs_layout, tuple(value.shape), lhs_shape
352 )
353 ndim_diff = len(lhs_shape) - value.ndim
354 if ndim_diff > 0:
355 padded_value = value.reshape((1,) * ndim_diff + tuple(value.shape))
356 else:
357 padded_value = value
358 local_value = _get_slice_tensor_by_layout(padded_value, value_layout)
359 return local_value, value_desc
360 return value, value_desc
362 return value, None
364 def infer_layout(self, cache_values: list) -> Tuple[tuple, None]: # pylint: disable=W0221
365 """Infer output layout for __setitem__.
367 __setitem__ is in-place, so output layout equals self layout.
368 Validates LHS key constraints and RHS value compatibility.
370 Rules:
371 1. self must not be in Partial status.
372 2. BoolTensor mask indexing is not supported.
373 3. Any dimension being written must be replicated.
374 4. slice step must be None or 1.
375 5. Advanced index tensors must be replicated.
376 6. RHS DTensor must have mesh consistent with self and
377 broadcast-compatible layout (non-broadcast dims match LHS
378 layout, broadcast dims must be Replicate).
379 7. RHS Tensor/scalar must be broadcastable to LHS slice shape.
380 8. Output layout equals self_layout (in-place).
382 Args:
383 cache_values: [self_layout, key_desc, global_shape, kind, value_desc, lhs_shape]
385 Returns:
386 tuple: ((self_layout,), None)
388 Raises:
389 ValueError: If any constraint is violated.
390 """
391 self_layout = cache_values[0]
392 key_desc = cache_values[1]
393 global_shape = cache_values[2]
394 kind = cache_values[3]
395 value_desc = cache_values[4]
396 lhs_shape = cache_values[5]
398 # key_desc -> expanded_actions (for LHS validation)
399 expanded_actions = _descriptor_to_expanded_actions(
400 key_desc, len(global_shape), op_name=self.op_name
401 )
403 if not self._allow_partial_inputs:
404 self._check_partial_inputs([self_layout])
405 self._validate_input_layouts(
406 self_layout, expanded_actions, global_shape, kind, value_desc, lhs_shape
407 )
408 return ((self_layout,), None)
410 @staticmethod
411 def _validate_input_layouts(self_layout, expanded_actions, global_shape, kind, value_desc, lhs_shape):
412 """Validate sharding constraints for __setitem__."""
413 op_name = "__setitem__"
415 GetItemDistributedOp._validate_input_layouts( # pylint: disable=W0212
416 self_layout, expanded_actions, global_shape, kind, op_name=op_name
417 )
419 if value_desc is None:
420 return
422 val_kind = value_desc[0]
423 if val_kind == "plain_tensor":
424 _validate_value_broadcast(op_name, value_desc[1], lhs_shape)
425 elif val_kind == "dtensor":
426 SetItemDistributedOp._validate_dtensor_value(
427 self_layout, expanded_actions, kind, value_desc[1], value_desc[2], lhs_shape, op_name,
428 )
429 else:
430 raise ValueError(
431 f"For {op_name}, unsupported value descriptor kind: {val_kind}."
432 )
434 @staticmethod
435 def _validate_dtensor_value(
436 self_layout, expanded_actions, kind, val_layout, val_shape, lhs_shape, op_name):
437 """Validate DTensor RHS mesh and layout."""
438 _validate_value_broadcast(op_name, val_shape, lhs_shape)
440 if val_layout.mesh.to_hash() != self_layout.mesh.to_hash():
441 raise ValueError(
442 f"For {op_name}, value DTensor must be on same mesh as self, "
443 f"but got mesh {val_layout.mesh_shape} "
444 f"vs {self_layout.mesh_shape}."
445 )
447 expected_layout = _compute_lhs_layout(
448 self_layout, expanded_actions, kind, op_name=op_name
449 )
450 SetItemDistributedOp._validate_value_layout(
451 val_layout, expected_layout, val_shape, lhs_shape, op_name,
452 )
454 @staticmethod
455 def _validate_value_layout(val_layout, expected_layout, val_shape, lhs_shape, op_name):
456 """Validate RHS DTensor placement against expected LHS placement."""
457 ndim_diff = len(lhs_shape) - len(val_shape)
458 val_alias = val_layout.alias_tensor_map
459 expected_alias = expected_layout.alias_tensor_map
461 for i, v_size in enumerate(val_shape):
462 lhs_index = ndim_diff + i
463 e_size = lhs_shape[lhs_index]
464 v_alias = val_alias[i]
465 e_alias = expected_alias[lhs_index]
467 if v_size == 1 and e_size > 1:
468 if v_alias != "None":
469 raise ValueError(
470 f"For {op_name}, value broadcast dim {i} (size {v_size}) "
471 f"must be Replicate, but got {v_alias}."
472 )
473 elif v_alias != e_alias:
474 raise ValueError(
475 f"For {op_name}, value layout mismatch at dim {i}: "
476 f"expected {e_alias}, but got {v_alias}."
477 )
479 def wrap_output(self, py_output: Any, output_layouts: tuple) -> None:
480 """Override wrap_output for __setitem__.
482 __setitem__ is in-place and returns None in PyTorch.
483 The default wrap_output would try to wrap None into a DTensor and fail.
485 Args:
486 py_output: The output from the local function call (None).
487 output_layouts: Output layouts from infer_layout.
489 Returns:
490 None: __setitem__ has no return value.
491 """
492 return None
495def _validate_value_broadcast(op_name, val_shape, lhs_shape):
496 """Validate that value shape can broadcast to LHS slice shape.
498 Args:
499 op_name: Operator name for error messages.
500 val_shape: Shape of RHS value.
501 lhs_shape: Expected LHS slice shape.
503 Raises:
504 ValueError: If shapes are not broadcastable.
505 """
506 if len(val_shape) > len(lhs_shape):
507 raise ValueError(
508 f"For {op_name}, value shape {val_shape} cannot broadcast "
509 f"to target shape {lhs_shape}."
510 )
512 # Pad val_shape with leading 1s
513 padded = (1,) * (len(lhs_shape) - len(val_shape)) + val_shape
514 for v, t in zip(padded, lhs_shape):
515 if v not in (1, t):
516 raise ValueError(
517 f"For {op_name}, value shape {val_shape} cannot broadcast "
518 f"to target shape {lhs_shape}."
519 )