Coverage for /home/jenkins/.local/lib/python3.10/site-packages/hyper_parallel/core/shard/ops/parallel

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

7# http://www.apache.org/licenses/LICENSE-2.0

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 __getitem__ operator (PyTorch only).

18Supports basic indexing (int, slice, None, Ellipsis) which produces views,

19and advanced indexing (list, LongTensor) which produces copies.

20BoolTensor masks are not supported because the output shape is data-dependent.

21"""

22from typing import Tuple

24from hyper_parallel.platform import get_platform

25from hyper_parallel.core.dtensor.dtensor import DTensor

26from hyper_parallel.core.dtensor.layout import Layout

27from .parallel_ops import DistributedOp

29platform = get_platform()

30Tensor = platform.Tensor

32_BASIC = "basic"

33_ADVANCED = "advanced"

34_BOOL_MASK = "bool_mask"

37def _normalize___getitem___args(self_t, key):

38 """Normalize __getitem__ arguments to canonical positional form.

40 __getitem__ is always called as func(self, key) with empty kwargs.

42 Args:

43 self_t: The DTensor self argument.

44 key: The indexing key (int, slice, None, Ellipsis, tuple, list, Tensor).

46 Returns:

47 tuple: ((self_t, key), {})

48 """

49 return (self_t, key), {}

52def _is_long_tensor(obj) -> bool:

53 """Check if obj is a non-bool Tensor (LongTensor or similar)."""

54 return isinstance(obj, Tensor) and not _is_bool_tensor(obj)

57def _is_bool_tensor(obj) -> bool:

58 """Check if obj is a BoolTensor."""

59 if not isinstance(obj, Tensor):

60 return False

61 if not hasattr(obj, 'dtype'):

62 return False

63 return str(obj.dtype).rsplit('.', maxsplit=1)[-1] in ('bool', 'bool_')

66def _is_advanced_elem(k) -> bool:

67 """Check if key element triggers advanced indexing."""

68 if isinstance(k, list):

69 return True

70 if isinstance(k, Tensor):

71 # 0-D long tensor is treated as basic (equivalent to int)

72 if k.ndim == 0 and not _is_bool_tensor(k):

73 return False

74 return True

75 return False

78def _build_key_element_descriptor(k, op_name="__getitem__"):

79 """Build a descriptor tuple for a single key element.

81 Args:

82 k: A single key element (int, slice, None, Ellipsis, list, Tensor).

83 op_name: Operator name for error messages.

85 Returns:

86 tuple: A descriptor tuple for this element.

88 Raises:

89 ValueError: If the key element type is unsupported.

90 """

91 if k is None:

92 return ("none",)

93 if k is Ellipsis:

94 return ("ellipsis",)

95 if isinstance(k, int):

96 return ("int", k)

97 if isinstance(k, slice):

98 return ("slice", k.start, k.stop, k.step)

99 if isinstance(k, list):

100 return ("idx_list_len", len(k))

101 if isinstance(k, Tensor) and k.ndim == 0 and not _is_bool_tensor(k):

102 return ("int", int(k.item()))

103 if isinstance(k, Tensor):

104 shape = tuple(k.shape)

105 if isinstance(k, DTensor):

106 alias = tuple(k.layout.alias_tensor_map)

107 return ("idx_tensor", shape, alias)

108 return ("idx_tensor", shape)

109 raise ValueError(

110 f"For {op_name}, unsupported index type: {type(k).__name__}."

111 )

112

113

114def _key_cache_descriptor(key, op_name="__getitem__"):

115 """Convert raw key to hashable (desc, kind) for cache_values.

116

117 The descriptor captures only what affects layout derivation:

118 int value, slice bounds, list length, and tensor shape.

119 Exact list/tensor contents are not preserved — only shape matters for layout.

120

121 Args:

122 key: Raw key from __getitem__ call (int, slice, None, Ellipsis, tuple, list, Tensor).

123 op_name: Operator name for error messages.

124

125 Returns:

126 (desc, kind): desc is a tuple of hashable descriptor tuples,

127 kind is "basic", "advanced", or "bool_mask".

128

129 Raises:

130 ValueError: If key contains unsupported types.

131 """

132 if not isinstance(key, tuple):

133 key = (key,)

134

135 # First pass: detect BoolTensor anywhere in the key

136 for k in key:

137 if _is_bool_tensor(k):

138 mask_shape = tuple(k.shape) if hasattr(k, 'shape') else ()

139 return (("bool_mask", mask_shape),), _BOOL_MASK

140

141 # Check for advanced indexing elements (list or non-0D tensor)

142 has_advanced = any(_is_advanced_elem(k) for k in key)

143 kind = _ADVANCED if has_advanced else _BASIC

144

145 desc = [_build_key_element_descriptor(k, op_name) for k in key]

146 return tuple(desc), kind

147

148

149def _desc_action_to_expanded(action_type, d, input_dim):

150 """Convert a single descriptor element to expanded action.

151

152 Args:

153 action_type: Type string from descriptor (none, int, slice, etc.).

154 d: Descriptor tuple for this element.

155 input_dim: Current input dimension index.

156

157 Returns:

158 tuple: (expanded_entries, input_dim_delta)

159 expanded_entries: list of action tuples to append.

160 input_dim_delta: how much to advance input_dim.

161 """

162 if action_type == "none":

163 return [("newaxis",)], 0

164 if action_type == "int":

165 return [("int", d[1], input_dim)], 1

166 if action_type == "slice":

167 return [("slice", d[1], d[2], d[3], input_dim)], 1

168 if action_type == "idx_list_len":

169 return [("idx_list", (0,) * d[1], input_dim)], 1

170 if action_type == "idx_tensor":

171 shape = d[1]

172 alias = d[2] if len(d) > 2 else None

173 idx_info = {"shape": shape, "alias": alias}

174 return [("idx_tensor", idx_info, input_dim)], 1

175 raise ValueError(

176 f"Unsupported descriptor type: {action_type}."

177 )

178

179

180def _descriptor_to_expanded_actions(desc, ndim, op_name="__getitem__"):

181 """Reconstruct expanded_actions from key descriptor for layout derivation.

182

183 Args:

184 desc: Key descriptor tuple from _key_cache_descriptor.

185 ndim: Number of input tensor dimensions.

186 op_name: Operator name for error messages.

187

188 Returns:

189 list: Expanded actions, each a tuple (action_type, *args).

190

191 Raises:

192 ValueError: On invalid descriptor (multiple Ellipsis, too many indices).

193 """

194 if desc is None:

195 return []

196

197 # Pass through bool_mask descriptor

198 if len(desc) == 1 and desc[0][0] == "bool_mask":

199 return [("bool_mask", desc[0][1])]

200

201 non_none = [d for d in desc if d[0] != "none"]

202 n_ellipsis = sum(1 for d in non_none if d[0] == "ellipsis")

203

204 if n_ellipsis > 1:

205 raise ValueError(

206 f"For {op_name}, an index can only have a single Ellipsis ('...')."

207 )

208

209 n_specified_dims = len(non_none) - n_ellipsis

210 n_fill = ndim - n_specified_dims

211

212 if n_fill < 0:

213 raise ValueError(

214 f"For {op_name}, too many indices for tensor of dimension {ndim}, "

215 f"but got {len(non_none)} non-None indices."

216 )

217

218 expanded = []

219 input_dim = 0

220

221 for d in desc:

222 action_type = d[0]

223

224 if action_type == "ellipsis":

225 for _ in range(n_fill):

226 expanded.append(("slice", None, None, None, input_dim))

227 input_dim += 1

228 else:

229 entries, dim_delta = _desc_action_to_expanded(action_type, d, input_dim)

230 expanded.extend(entries)

231 input_dim += dim_delta

232

233 # Implicitly add full slices for unspecified trailing dimensions

234 while input_dim < ndim:

235 expanded.append(("slice", None, None, None, input_dim))

236 input_dim += 1

237

238 return expanded

239

240

241def _unwrap_key_for_local(key):

242 """Convert key for local execution: DTensors to local.

243

244 Args:

245 key: Raw key from __getitem__ call.

246

247 Returns:

248 Unwrapped key suitable for passing to local func.

249 """

250 if isinstance(key, DTensor):

251 return key.to_local()

252 if isinstance(key, tuple):

253 return tuple(_unwrap_key_for_local(k) for k in key)

254 return key

255

256

257def _copy_partial_state(src_layout, dst_layout):

258 """Copy partial state from src_layout to dst_layout using public API.

259

260 Args:

261 src_layout: Source Layout to copy partials from.

262 dst_layout: Destination Layout to copy partials to.

263 """

264 for dev_idx, op in enumerate(src_layout.partial):

265 if op is not None:

266 dst_layout.set_partial_by_dev_axis(dst_layout.alias_name[dev_idx], op)

267

268

269

270class GetItemDistributedOp(DistributedOp):

271 """Distributed implementation for tensor.__getitem__.

272

273 Supports basic indexing (int, slice, None, Ellipsis) which produces views,

274 and advanced indexing (list, LongTensor) which produces copies.

275 BoolTensor masks are rejected because they produce data-dependent shapes.

276

277 Sharding constraints:

278 - Any dimension indexed by int, non-full slice, or advanced index must

279 be replicated.

280 - Advanced index tensors must themselves be replicated.

281 - Input must not have Partial status.

282 - slice step != 1 is not supported.

283 """

284

285 def preprocess(self, args: tuple, kwargs: dict) -> tuple:

286 """Preprocess arguments for __getitem__.

287

288 Converts raw key to hashable descriptor for cache, and unwraps

289 DTensors to local tensors for execution.

290

291 Args:

292 args: (self_tensor, key)

293 kwargs: Empty dict for __getitem__.

294

295 Returns:

296 tuple: (local_args, local_kwargs, cache_values)

297 """

298 norm_args, _ = _normalize___getitem___args(*args, **kwargs)

299 self_t, key = norm_args

300

301 self_layout = self_t.layout

302 global_shape = tuple(self_t.shape)

303

304 # raw key -> hashable descriptor for cache

305 key_desc, kind = _key_cache_descriptor(key, op_name=self.op_name)

306

307 # Unwrap DTensor in key to local for local execution

308 local_key = _unwrap_key_for_local(key)

309 local_args = (self_t.to_local(), local_key)

310 local_kwargs = {}

311

312 cache_values = [self_layout, key_desc, global_shape, kind]

313 return local_args, local_kwargs, cache_values

314

315 @staticmethod

316 def _reject_bool_mask(expanded_actions, op_name="__getitem__"):

317 """Raise ValueError for bool_mask key kind."""

318 mask_shape = "unknown"

319 for action in expanded_actions:

320 if action[0] == "bool_mask":

321 mask_shape = action[1]

322 break

323 raise ValueError(

324 f"For {op_name}, boolean-mask indexing has data-dependent "

325 f"output shape and is not supported in DTensor. "

326 f"Got mask of shape {mask_shape}."

327 )

328

329 @staticmethod

330 def _validate_int_action(action, alias_map, global_shape, op_name="__getitem__"):

331 """Validate an int indexing action."""

332 input_dim = action[-1]

333 idx = action[1]

334 if idx < -global_shape[input_dim] or idx >= global_shape[input_dim]:

335 raise ValueError(

336 f"For {op_name}, index {idx} is out of range "

337 f"for dimension {input_dim} with size {global_shape[input_dim]}."

338 )

339 if alias_map[input_dim] != "None":

340 raise ValueError(

341 f"For {op_name}, indexing with int on non-replicate "

342 f"dim {input_dim} is not supported, "

343 f"but got sharding {alias_map[input_dim]} on dim {input_dim}."

344 )

345

346 @staticmethod

347 def _validate_slice_action(action, alias_map, global_shape, op_name="__getitem__"):

348 """Validate a slice indexing action."""

349 input_dim = action[-1]

350 start, stop, step = action[1], action[2], action[3]

351 step = step if step is not None else 1

352

353 if step != 1:

354 raise ValueError(

355 f"For {op_name}, slice step should be 1 or None, "

356 f"but got {step}."

357 )

358

359 is_full = (start is None or start == 0) and (

360 stop is None or stop >= global_shape[input_dim]

361 )

362

363 if not is_full and alias_map[input_dim] != "None":

364 raise ValueError(

365 f"For {op_name}, non-full slice on non-replicate "

366 f"dim {input_dim} is not supported, "

367 f"but got sharding {alias_map[input_dim]} on dim {input_dim}."

368 )

369

370 @staticmethod

371 def _validate_advanced_action(action, alias_map, op_name="__getitem__"):

372 """Validate an advanced indexing action (idx_list or idx_tensor)."""

373 input_dim = action[-1]

374 if alias_map[input_dim] != "None":

375 raise ValueError(

376 f"For {op_name}, advanced indexing on non-replicate "

377 f"dim {input_dim} is not supported, "

378 f"but got sharding {alias_map[input_dim]} on dim {input_dim}."

379 )

380 # For idx_tensor, the index tensor itself must be replicated

381 if action[0] == "idx_tensor":

382 idx_alias = action[1].get("alias")

383 if idx_alias and any(x != "None" for x in idx_alias):

384 raise ValueError(

385 f"For {op_name}, advanced index tensor must be "

386 f"replicated, but got layout with sharding {idx_alias}."

387 )

388

389 @staticmethod

390 def _validate_input_layouts(self_layout, expanded_actions, global_shape, kind,

391 op_name="__getitem__"):

392 """Validate sharding constraints for __getitem__.

393

394 Rules:

395 1. BoolTensor mask indexing is not supported.

396 2. Any dimension indexed by int, non-full slice, or advanced index

397 must be replicated.

398 3. slice step must be None or 1.

399 4. Advanced index tensors must be replicated.

400 5. int indices must be in range.

401

402 Args:

403 self_layout: Layout of self tensor.

404 expanded_actions: Expanded key actions from _descriptor_to_expanded_actions.

405 global_shape: Global shape of self tensor.

406 kind: "basic" or "advanced".

407 op_name: Operator name for error messages (default "__getitem__").

408

409 Raises:

410 ValueError: If any constraint is violated.

411 """

412 alias_map = self_layout.alias_tensor_map

413

414 if kind == _BOOL_MASK:

415 GetItemDistributedOp._reject_bool_mask(expanded_actions, op_name)

416

417 for action in expanded_actions:

418 action_type = action[0]

419

420 if action_type == "newaxis":

421 continue

422 if action_type == "int":

423 GetItemDistributedOp._validate_int_action(

424 action, alias_map, global_shape, op_name

425 )

426 elif action_type == "slice":

427 GetItemDistributedOp._validate_slice_action(

428 action, alias_map, global_shape, op_name

429 )

430 elif action_type in ("idx_list", "idx_tensor"):

431 GetItemDistributedOp._validate_advanced_action(

432 action, alias_map, op_name

433 )

434

435 def infer_layout(self, cache_values: list) -> Tuple[tuple, None]: # pylint: disable=W0221

436 """Infer output layout for __getitem__.

437

438 Rules:

439 1. Input must not have Partial status.

440 2. BoolTensor mask indexing is not supported.

441 3. Any dimension indexed by int, non-full slice, or advanced index

442 must be replicated.

443 4. slice step must be None or 1.

444 5. Advanced index tensors must be replicated.

445 6. BASIC: Output alias_tensor_map is derived by removing int-indexed

446 dims, inserting Replicate for newaxis, and preserving sharding

447 for full-slice dims.

448 7. ADVANCED: Advanced indices broadcast shape B is inserted at

449 position p (consecutive: p = first advanced dim; non-consecutive:

450 p = 0). B dims are all Replicate. Other dims preserve sharding.

451 8. Partial state is copied to the output layout for preserved dims.

452

453 Args:

454 cache_values: [self_layout, key_desc, global_shape, kind]

455

456 Returns:

457 tuple: ((output_layout,), None)

458

459 Raises:

460 ValueError: If any constraint is violated.

461 """

462 self_layout = cache_values[0]

463 key_desc = cache_values[1]

464 global_shape = cache_values[2]

465 kind = cache_values[3]

466

467 # key_desc -> expanded_actions

468 expanded_actions = _descriptor_to_expanded_actions(

469 key_desc, len(global_shape), op_name=self.op_name

470 )

471

472 if not self._allow_partial_inputs:

473 self._check_partial_inputs([self_layout])

474 self._validate_input_layouts(self_layout, expanded_actions, global_shape, kind)

475

476 if kind == _BASIC:

477 out_layout = self._infer_basic_output_layout(self_layout, expanded_actions)

478 elif kind == _ADVANCED:

479 out_layout = self._infer_advanced_output_layout(

480 self_layout, expanded_actions, op_name=self.op_name

481 )

482 else:

483 # _BOOL_MASK is rejected by _validate_input_layouts above,

484 # so we should never reach this branch.

485 raise ValueError(

486 f"For __getitem__, unexpected kind: {kind}. "

487 f"Expected 'basic' or 'advanced'."

488 )

489

490 return ((out_layout,), None)

491

492 @staticmethod

493 def _infer_basic_output_layout(self_layout, expanded_actions):

494 """Derive output layout for basic indexing.

495

496 Walk through expanded_actions:

497 - int: remove the dimension from alias_tensor_map

498 - full slice (None, None, 1): keep the dimension with same sharding

499 - non-full slice: keep the dimension, must be replicated (validated earlier)

500 - newaxis: insert Replicate dimension

501

502 Args:

503 self_layout: Layout of self tensor.

504 expanded_actions: Expanded key actions.

505

506 Returns:

507 Layout: Output layout.

508 """

509 alias_map = self_layout.alias_tensor_map

510 mesh = self_layout.mesh

511

512 out_alias = []

513 for action in expanded_actions:

514 action_type = action[0]

515 if action_type == "newaxis":

516 out_alias.append("None")

517 elif action_type == "slice":

518 input_dim = action[4]

519 out_alias.append(alias_map[input_dim])

520 # "int" actions are skipped (dimension removed)

521

522 out_layout = Layout.from_device_mesh(mesh)

523 out_layout = out_layout(*out_alias)

524

525 # Copy partial state from input

526 _copy_partial_state(self_layout, out_layout)

527 out_layout.tensor_map_to_placement()

528 out_layout.update_compact_str()

529 return out_layout

530

531 @staticmethod

532 def _compute_advanced_index_info(expanded_actions, op_name):

533 """Analyze advanced indexing actions for output layout derivation.

534

535 Returns:

536 (advanced_positions, advanced_pos_set, are_consecutive, b_ndim)

537 """

538 advanced_actions = []

539 for i, action in enumerate(expanded_actions):

540 if action[0] in ("idx_list", "idx_tensor"):

541 advanced_actions.append((i, action))

542

543 index_shapes = []

544 for _, action in advanced_actions:

545 if action[0] == "idx_list":

546 index_shapes.append((len(action[1]),))

547 elif action[0] == "idx_tensor":

548 index_shapes.append(tuple(action[1]["shape"]))

549

550 bcast_shape = _broadcast_shapes(index_shapes, op_name=op_name)

551 b_ndim = len(bcast_shape) if bcast_shape else 0

552

553 advanced_positions = [pos for pos, _ in advanced_actions]

554 advanced_pos_set = set(advanced_positions)

555 are_consecutive = (

556 len(advanced_positions) > 0

557 and advanced_positions == list(range(advanced_positions[0],

558 advanced_positions[-1] + 1))

559 )

560 return advanced_positions, advanced_pos_set, are_consecutive, b_ndim

561

562 @staticmethod

563 def _infer_advanced_output_layout(self_layout, expanded_actions, op_name="__getitem__"):

564 """Derive output layout for advanced indexing.

565

566 Advanced indices are on specific input dims L. Other input dims K

567 (including newaxis) are preserved. The broadcast shape B of the

568 index tensors is inserted at position p:

569 - consecutive advanced indices: p = position of first advanced action

570 - non-consecutive: p = 0 (B precedes all other dims)

571

572 Args:

573 self_layout: Layout of self tensor.

574 expanded_actions: Expanded key actions.

575 op_name: Operator name for error messages.

576

577 Returns:

578 Layout: Output layout.

579 """

580 alias_map = self_layout.alias_tensor_map

581 mesh = self_layout.mesh

582

583 advanced_positions, advanced_pos_set, are_consecutive, b_ndim = \

584 GetItemDistributedOp._compute_advanced_index_info(expanded_actions, op_name)

585

586 def _append_non_advanced(action):

587 """Append output alias for a non-advanced action (newaxis, slice, int)."""

588 if action[0] == "newaxis":

589 out_alias.append("None")

590 elif action[0] == "slice":

591 out_alias.append(alias_map[action[4]])

592 # int: dimension removed, skip

593

594 out_alias = []

595

596 if are_consecutive and advanced_positions:

597 first_adv_pos = advanced_positions[0]

598 last_adv_pos = advanced_positions[-1]

599

600 for i in range(first_adv_pos):

601 _append_non_advanced(expanded_actions[i])

602 for _ in range(b_ndim):

603 out_alias.append("None")

604 for i in range(last_adv_pos + 1, len(expanded_actions)):

605 _append_non_advanced(expanded_actions[i])

606 else:

607 for _ in range(b_ndim):

608 out_alias.append("None")

609 for i, action in enumerate(expanded_actions):

610 if i in advanced_pos_set:

611 continue

612 _append_non_advanced(action)

613

614 out_layout = Layout.from_device_mesh(mesh)

615 out_layout = out_layout(*out_alias)

616

617 _copy_partial_state(self_layout, out_layout)

618 out_layout.tensor_map_to_placement()

619 out_layout.update_compact_str()

620 return out_layout

621

622

623

624def _broadcast_shapes(shapes, op_name="__getitem__"):

625 """Compute broadcast shape for a list of shapes.

626

627 Args:

628 shapes: List of shape tuples.

629 op_name: Operator name for error messages.

630

631 Returns:

632 tuple: Broadcast shape.

633 """

634 if not shapes:

635 return ()

636 result = list(shapes[0])

637 for shape in shapes[1:]:

638 ndim_diff = len(result) - len(shape)

639 if ndim_diff < 0:

640 result = [1] * (-ndim_diff) + result

641 ndim_diff = 0

642 for i, d2 in enumerate(shape):

643 d1 = result[ndim_diff + i]

644 if d1 == 1:

645 result[ndim_diff + i] = d2

646 elif d2 not in (1, d1):

647 raise ValueError(

648 f"For {op_name}, advanced index shapes {shapes} "

649 f"cannot be broadcast together."

650 )

651 return tuple(result)

Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / shard / ops / parallel_getitem.py: 90%

273 statements