Coverage for /home/jenkins/.local/lib/python3.10/site-packages/hyper_parallel/core/dtensor/redistribute

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"""redistribute_infer"""

16from typing import Dict, List, Tuple, Union

19class Status:

20 SUCCESS = 0

21 FAILED = 1

24CONCAT_BY_AXIS = 0

25SPLIT_BY_AXIS = 1

26PERMUTE_BY_AXIS = 2

27NONE = -1

30class TensorMap:

31 """Enhanced tensor map struct supporting tuples for combined dimensions"""

32 def __init__(self, dims: List[Union[int, Tuple[int, ...]]]):

33 self.dims = dims

35 def get_dim_by_idx(self, index: int) -> Union[int, Tuple[int, ...]]:

36 """Return the dimension value at the given index, or NONE if out of range."""

37 return self.dims[index] if index < len(self.dims) else NONE

39 def get_index_by_value(self, value: Union[int, Tuple[int, ...]]) -> int:

40 """Return the index of the first dimension matching the given value, or NONE."""

41 for i, dim in enumerate(self.dims):

42 if dim == value:

43 return i

44 return NONE

46 def get_index_contain_value(self, value: Union[int, Tuple[int, ...]]) -> int:

47 """Return the index of the tuple dimension whose suffix matches the given value, or NONE."""

48 for i, dim in enumerate(self.dims):

49 if not isinstance(dim, tuple):

50 continue

51 if isinstance(value, tuple) and value == dim[len(dim) - len(value):]:

52 return i

53 if not isinstance(value, tuple) and value == dim[-1]:

54 return i

55 return NONE

58class DevMat:

59 """

60 Represents a multi-dimensional grid of devices where each dimension has a specific size.

61 Supports operations to retrieve device groups along single or combined dimensions.

63 Attributes:

64 dims (List[int]): Sizes of each dimension in the mesh shape.

65 _combined_dims (Dict[Tuple[int, ...], int]): Cache for precomputed combined dimension sizes.

66 """

68 def __init__(self, dims: List[int]):

69 """

70 Initialize mesh shape dimensions.

72 Args:

73 dims: List of integers representing the size of each dimension.

74 """

75 self.dims = dims

76 self._combined_dims: Dict[Tuple[int, ...], int] = {}

78 def get_dim_by_reverse_idx(self, idx: Union[int, Tuple[int, ...]]) -> int:

79 """

80 Get dimension size by reverse index or product of combined dimensions.

82 For a single integer index `i`, returns the size of the dimension at reverse

83 position (i.e., `dims[len(dims)-1-i]`). For a tuple of indices, returns the

84 product of sizes for the specified reverse-indexed dimensions.

86 Args:

87 idx: Integer dimension index or tuple of indices.

89 Returns:

90 Dimension size (for integer) or product of sizes (for tuple).

91 """

92 if isinstance(idx, tuple):

93 return self._get_combined_size(idx)

94 return self.dims[len(self.dims) - 1 - idx]

96 def _get_combined_size(self, dims: Union[int, Tuple[int, ...]]) -> int:

97 """

98 Compute and cache the product of sizes for combined dimensions.

100 Args:

101 dims: Tuple of dimension indices (reverse-indexed).

102

103 Returns:

104 Product of sizes for the specified dimensions.

105 """

106 if dims in self._combined_dims:

107 return self._combined_dims[dims]

108 size = 1

109 for d in dims:

110 size *= self.dims[len(self.dims) - 1 - d]

111 self._combined_dims[dims] = size

112 return size

113

114 def _get_devices_along_dim(self, rank: int, rank_list: List[int], dim: int) -> List[int]:

115 """

116 Get devices sharing the same coordinates.

117

118 Devices are grouped such that only the specified dimension varies. The mesh shape

119 is assumed to be in row-major order (last dimension changes fastest).

120

121 Args:

122 rank: Target device rank.

123 rank_list: Flattened list of all devices in row-major order.

124 dim: Target dimension index (0-indexed from outermost).

125

126 Returns:

127 List of devices in the same group as `rank` along `dim`.

128

129 Raises:

130 ValueError: For invalid dimension or mismatched rank_list size.

131 """

132 if dim < 0 or dim >= len(self.dims):

133 raise ValueError(f"Dimension {dim} out of range [0, {len(self.dims)})")

134

135 # Trivial case: dimension size is 1

136 if self.dims[dim] == 1:

137 return [rank]

138

139 total_devices = 1

140 for d in self.dims:

141 total_devices *= d

142

143 # Validate rank_list length

144 if len(rank_list) != total_devices:

145 raise ValueError(f"rank_list length ({len(rank_list)}) doesn't match "

146 f"mesh shape product ({total_devices})")

147

148 # Compute stride for the dimension

149 stride = 1

150 for i in range(dim + 1, len(self.dims)):

151 stride *= self.dims[i]

152

153 # Find local index of rank in rank_list

154 try:

155 local_index = rank_list.index(rank)

156 except ValueError as e:

157 raise ValueError(f"Rank {rank} not in rank_list") from e

158

159 # Calculate base index and generate group

160 index_in_dim = (local_index // stride) % self.dims[dim]

161 base = local_index - index_in_dim * stride

162 group = [rank_list[base + k * stride] for k in range(self.dims[dim])]

163

164 return group

165

166 def get_devices_along_dim(self, rank: int, rank_list: List[int], dim: Union[int, List[int]]) -> List[int]:

167 """

168 Get devices sharing the same coordinates.

169

170 For a single dimension, returns devices where only that dimension varies.

171 For a tuple of dimensions, returns devices where ONLY the specified dimensions vary,

172 sharing fixed coordinates in all other dimensions.

173

174 Args:

175 rank: Target device rank.

176 rank_list: Flattened list of all devices in row-major order.

177 dim: Single dimension index or tuple of indices.

178

179 Returns:

180 List of devices in the same hyperplane as `rank` orthogonal to `dim`.

181

182 Raises:

183 ValueError: For invalid dimensions or mismatched rank_list size.

184 """

185 if isinstance(dim, list):

186 result = self._get_devices_along_dim(rank, rank_list, dim[0])

187 current_layer_len = len(result)

188 current_layer_step = 0

189 dim_index = 1

190 while dim_index < len(dim):

191 sub_rank = result.pop(0)

192 result.extend(self._get_devices_along_dim(sub_rank, rank_list, dim[dim_index]))

193 current_layer_step += 1

194 if current_layer_step == current_layer_len:

195 dim_index += 1

196 current_layer_step = 0

197 current_layer_len = len(result)

198 return result

199 return self._get_devices_along_dim(rank, rank_list, dim)

200

201

202class RedistributionOperatorInfer:

203 """

204 Infers communication operators for tensor redistribution in distributed systems.

205

206 Determines the sequence of communication operations (split, concat, permute)

207 required to transform a tensor from an input device mapping to an output device mapping.

208

209 Args:

210 dev_mat: Mesh shape dimensions representing the device grid

211 in_tensor_map: Input tensor's device mapping for each tensor dimension

212 out_tensor_map: Output tensor's device mapping for each tensor dimension

213 use_permute: Whether to use permute operator (all-to-all) when possible (default: True)

214 """

215 def __init__(self, dev_mat: List[int],

216 in_tensor_map: List[Union[int, Tuple[int, ...]]],

217 out_tensor_map: List[Union[int, Tuple[int, ...]]],

218 use_permute: bool = True):

219

220 self.operator_list_: List[Tuple[int, Tuple]] = []

221 self.map_: Dict[int, Union[int, Tuple[int, ...]]] = {}

222 self.use_permute = use_permute

223

224 # Initialize with expanded dimensions

225 self.dev_ranks = len(dev_mat)

226 self.dev_mat_ = DevMat(dev_mat)

227 self.in_tensor_map_ = TensorMap(in_tensor_map)

228 self.out_tensor_map_ = TensorMap(out_tensor_map)

229

230 self.map_ = {i: self.in_tensor_map_.get_dim_by_idx(i)

231 for i in range(len(in_tensor_map))}

232

233 def insert_operator(self, op_type: int, args: Tuple) -> int:

234 """

235 Adds an operator to the internal operator sequence.

236

237 Args:

238 op_type: Operator type constant (SPLIT_BY_AXIS, CONCAT_BY_AXIS, PERMUTE_BY_AXIS)

239 args: Operator-specific arguments tuple

240

241 Returns:

242 Status.SUCCESS on success, Status.FAILED on error

243 """

244 self.operator_list_.append((op_type, args))

245 return Status.SUCCESS

246

247 def infer_redistribution_operator(self) -> int:

248 """

249 Main inference driver coordinating the redistribution sequence.

250

251 Executes in 3 phases until mapping is resolved:

252 1. Split operations

253 2. Permute/All-to-All operations

254 3. Concat operations

255

256 Returns:

257 Status.SUCCESS if full sequence inferred, Status.FAILED otherwise

258 """

259 while self.map_:

260 len_global = len(self.operator_list_)

261

262 while self.map_:

263 len_split_by_axis = len(self.operator_list_)

264

265 # Step 1: infer split op

266 if self.infer_split_by_axis() == Status.FAILED:

267 return Status.FAILED

268

269 # Step 2: infer alltoall op

270 while self.map_:

271 len_permute_by_axis = len(self.operator_list_)

272 if self.infer_permute_by_axis() == Status.FAILED:

273 return Status.FAILED

274 if len_permute_by_axis == len(self.operator_list_):

275 break

276

277 if len_split_by_axis == len(self.operator_list_):

278 break

279

280 # Step 3: infer allconcat op

281 if self.infer_concat_by_axis() == Status.FAILED:

282 return Status.FAILED

283

284 if len_global == len(self.operator_list_) and self.map_:

285 index = next(iter(self.map_.keys()))

286 in_dim = self.map_[index]

287 self.map_[index] = NONE

288 dev_dim = self.dev_mat_.get_dim_by_reverse_idx(in_dim)

289 args = (index, in_dim, dev_dim)

290 if self.insert_operator(CONCAT_BY_AXIS, args) == Status.FAILED:

291 return Status.FAILED

292

293 return Status.SUCCESS

294

295 def _handle_simple_split_case(self, index: int, in_dim: Union[int, Tuple[int, ...]],

296 out_dim: Union[int, Tuple[int, ...]]) -> bool:

297 """Handle the simple case where input dimension is None and output dimension is not conflicting"""

298 if in_dim != NONE:

299 return False

300

301 conflict = any(v == out_dim for v in self.map_.values())

302 if isinstance(out_dim, tuple):

303 conflict_tuple = any(isinstance(v, tuple) and v[len(v) - len(out_dim):] == out_dim

304 for v in self.map_.values())

305 else:

306 conflict_tuple = any(isinstance(v, tuple) and v[-1] == out_dim for v in self.map_.values())

307

308 if not conflict and not conflict_tuple:

309 dev_dim = self.dev_mat_.get_dim_by_reverse_idx(out_dim)

310 args = (index, out_dim, dev_dim)

311 return self.insert_operator(SPLIT_BY_AXIS, args) == Status.SUCCESS

312

313 return False

314

315 def _handle_tuple_split_case(self, index: int, in_dim: Union[int, Tuple[int, ...]],

316 out_dim: Union[int, Tuple[int, ...]]) -> bool:

317 """Handle the case where output dimension is a tuple and input dimension matches prefix"""

318 if not isinstance(out_dim, tuple):

319 return False

320

321 in_dim_matches_out_prefix = (

322 (not isinstance(in_dim, tuple) and in_dim == out_dim[0]) or

323 (isinstance(in_dim, tuple) and in_dim == out_dim[:len(in_dim)])

324 )

325 if in_dim_matches_out_prefix:

326

327 if isinstance(in_dim, tuple):

328 out_dim_rest = out_dim[-1] if len(out_dim[len(in_dim):]) == 1 else out_dim[len(in_dim):]

329 else:

330 out_dim_rest = out_dim[-1] if len(out_dim[1:]) == 1 else out_dim[1:]

331

332 conflict = any(v == out_dim_rest for v in self.map_.values())

333 if not conflict:

334 dev_dim = self.dev_mat_.get_dim_by_reverse_idx(out_dim_rest)

335 args = (index, out_dim_rest, dev_dim)

336 return self.insert_operator(SPLIT_BY_AXIS, args) == Status.SUCCESS

337

338 return False

339

340 def infer_split_by_axis(self) -> int:

341 """

342 Infers split operations for the current mapping state.

343

344 Conditions for split:

345 - Tensor dimension changes from unmapped to mapped

346 - No conflicts in target device dimension

347

348 Updates internal mapping state and operator list.

349

350 Returns:

351 Status.SUCCESS if operations inferred, Status.FAILED on error

352 """

353 keys = list(self.map_.keys())

354 for index in keys:

355 if index not in self.map_:

356 continue

357

358 in_dim = self.map_[index]

359 out_dim = self.out_tensor_map_.get_dim_by_idx(index)

360

361 if in_dim == out_dim:

362 del self.map_[index]

363 continue

364

365 # Handle simple case: input dimension is None

366 if self._handle_simple_split_case(index, in_dim, out_dim):

367 del self.map_[index]

368 continue

369

370 # Handle tuple case: output dimension is a tuple

371 if self._handle_tuple_split_case(index, in_dim, out_dim):

372 del self.map_[index]

373 continue

374

375 return Status.SUCCESS

376

377 def _handle_none_dim_permute_case(self, index: int, in_dim: Union[int, Tuple[int, ...]],

378 out_dim: Union[int, Tuple[int, ...]]) -> bool:

379 """Handle permute case where input dimension is None"""

380 if in_dim != NONE:

381 return False

382

383 # Check for conflicts in output dimension

384 conflict = any(v == out_dim for v in self.map_.values())

385 if not conflict:

386 return False

387

388 # Handle regular dimension conflict

389 concat_axis = self.in_tensor_map_.get_index_by_value(out_dim)

390 if concat_axis is None:

391 return False

392

393 split_dev_num = self.dev_mat_.get_dim_by_reverse_idx(out_dim)

394

395 if self.use_permute:

396 # concat tensor map value, to get the communication group

397 concat_map = self.in_tensor_map_.get_dim_by_idx(concat_axis)

398 concat_dev_num = self.dev_mat_.get_dim_by_reverse_idx(concat_map)

399 args_permute = (concat_dev_num, index, concat_axis, concat_map, split_dev_num)

400

401 if self.insert_operator(PERMUTE_BY_AXIS, args_permute) == Status.FAILED:

402 return False

403 else:

404 args_concat = (concat_axis, out_dim, split_dev_num)

405 args_split = (index, out_dim, split_dev_num)

406

407 if self.insert_operator(CONCAT_BY_AXIS, args_concat) == Status.FAILED:

408 return False

409 if self.insert_operator(SPLIT_BY_AXIS, args_split) == Status.FAILED:

410 return False

411

412 del self.map_[index]

413 self.map_[concat_axis] = NONE

414 return True

415

416 def _handle_none_dim_tuple_permute_case(self, index: int, in_dim: Union[int, Tuple[int, ...]],

417 out_dim: Union[int, Tuple[int, ...]]) -> bool:

418 """Handle permute case where input dimension is None and output dimension is a tuple with conflicts"""

419 if in_dim != NONE:

420 return False

421

422 if isinstance(out_dim, tuple):

423 conflict_tuple = any(isinstance(v, tuple) and v[len(v) - len(out_dim):] == out_dim

424 for v in self.map_.values())

425 else:

426 conflict_tuple = any(isinstance(v, tuple) and v[-1] == out_dim for v in self.map_.values())

427

428 if not conflict_tuple:

429 return False

430

431 concat_axis = self.in_tensor_map_.get_index_contain_value(out_dim)

432 if concat_axis is None:

433 return False

434

435 split_dev_num = self.dev_mat_.get_dim_by_reverse_idx(out_dim)

436

437 if self.use_permute:

438 # concat tensor map value, to get the communication group

439 concat_map = out_dim

440 concat_dev_num = self.dev_mat_.get_dim_by_reverse_idx(concat_map)

441 args_permute = (concat_dev_num, index, concat_axis, concat_map, split_dev_num)

442

443 if self.insert_operator(PERMUTE_BY_AXIS, args_permute) == Status.FAILED:

444 return False

445 else:

446 args_concat = (concat_axis, out_dim, split_dev_num)

447 args_split = (index, out_dim, split_dev_num)

448

449 if self.insert_operator(CONCAT_BY_AXIS, args_concat) == Status.FAILED:

450 return False

451 if self.insert_operator(SPLIT_BY_AXIS, args_split) == Status.FAILED:

452 return False

453

454 del self.map_[index]

455 out_dim_len = 1 if not isinstance(out_dim, tuple) else len(out_dim)

456 rest_size = len(self.map_[concat_axis]) - out_dim_len

457 new_map_item = self.map_[concat_axis][:rest_size] if rest_size > 1 else self.map_[concat_axis][0]

458 self.map_[concat_axis] = new_map_item

459 return True

460

461 def _handle_tuple_dim_permute_case(self, index: int, in_dim: Union[int, Tuple[int, ...]],

462 out_dim: Union[int, Tuple[int, ...]]) -> bool:

463 """Handle permute case where both input and output dimensions are tuples"""

464 if not isinstance(out_dim, tuple):

465 return False

466

467 if not ((not isinstance(in_dim, tuple) and in_dim == out_dim[0]) or

468 (isinstance(in_dim, tuple) and in_dim == out_dim[:len(in_dim)])):

469 return False

470

471 if isinstance(in_dim, tuple):

472 out_dim_rest = out_dim[-1] if len(out_dim[len(in_dim):]) == 1 else out_dim[len(in_dim):]

473 else:

474 out_dim_rest = out_dim[-1] if len(out_dim[1:]) == 1 else out_dim[1:]

475

476 conflict = any(v == out_dim_rest for v in self.map_.values())

477 if not conflict:

478 return False

479

480 concat_axis = self.in_tensor_map_.get_index_by_value(out_dim_rest)

481 if concat_axis is None:

482 return False

483

484 split_dev_num = self.dev_mat_.get_dim_by_reverse_idx(out_dim_rest)

485

486 if self.use_permute:

487 # concat tensor map value, to get the communication group

488 concat_map = out_dim_rest

489 concat_dev_num = self.dev_mat_.get_dim_by_reverse_idx(concat_map)

490 args_permute = (concat_dev_num, index, concat_axis, concat_map, split_dev_num)

491

492 if self.insert_operator(PERMUTE_BY_AXIS, args_permute) == Status.FAILED:

493 return False

494 else:

495 args_concat = (concat_axis, out_dim_rest, split_dev_num)

496 args_split = (index, out_dim_rest, split_dev_num)

497

498 if self.insert_operator(CONCAT_BY_AXIS, args_concat) == Status.FAILED:

499 return False

500 if self.insert_operator(SPLIT_BY_AXIS, args_split) == Status.FAILED:

501 return False

502

503 del self.map_[index]

504 self.map_[concat_axis] = NONE

505 return True

506

507 def infer_permute_by_axis(self) -> int:

508 """

509 Infers permutation (all-to-all) operations for dimension conflicts.

510

511 Handles cases where:

512 - Input dimension is unmapped but output dimension is already occupied

513 - Uses either permute operator or split+concat pair based on use_permute flag

514

515 Returns:

516 Status.SUCCESS if operations inferred, Status.FAILED on error

517 """

518 keys = list(self.map_.keys())

519 for index in keys:

520 if index not in self.map_:

521 continue

522

523 in_dim = self.map_[index]

524 out_dim = self.out_tensor_map_.get_dim_by_idx(index)

525

526 if in_dim == out_dim:

527 del self.map_[index]

528 continue

529

530 # Handle different permute cases

531 if self._handle_none_dim_permute_case(index, in_dim, out_dim):

532 continue

533

534 if self._handle_none_dim_tuple_permute_case(index, in_dim, out_dim):

535 continue

536

537 if self._handle_tuple_dim_permute_case(index, in_dim, out_dim):

538 continue

539

540 return Status.SUCCESS

541

542 def _handle_tuple_concat_case(self, index: int, in_dim: Union[int, Tuple[int, ...]],

543 out_dim: Union[int, Tuple[int, ...]]) -> bool:

544 """Handle concat case where input dimension is a tuple and output matches prefix"""

545 if not isinstance(in_dim, tuple):

546 return False

547

548 if out_dim == NONE:

549 if len(in_dim) <= 1:

550 return False

551 # Plain same-dim Shard produces a descending tuple, so the original

552 # combined concat is enough. StridedShard reorders the tuple.

553 if all(in_dim[i] > in_dim[i + 1] for i in range(len(in_dim) - 1)):

554 return False

555 in_dim_rest = in_dim[-1]

556 concat_dev_num = self.dev_mat_.get_dim_by_reverse_idx(in_dim_rest)

557 args = (index, in_dim_rest, concat_dev_num)

558

559 if self.insert_operator(CONCAT_BY_AXIS, args) == Status.FAILED:

560 return False

561

562 self.map_[index] = in_dim[:-1] if len(in_dim) > 2 else in_dim[0]

563 return True

564

565 if not ((not isinstance(out_dim, tuple) and out_dim == in_dim[0]) or

566 (isinstance(out_dim, tuple) and out_dim == in_dim[:len(out_dim)])):

567 return False

568

569 if isinstance(out_dim, tuple):

570 in_dim_rest = in_dim[-1] if len(in_dim[len(out_dim):]) == 1 else in_dim[len(out_dim):]

571 else:

572 in_dim_rest = in_dim[-1] if len(in_dim[1:]) == 1 else in_dim[1:]

573

574 concat_dev_num = self.dev_mat_.get_dim_by_reverse_idx(in_dim_rest)

575 args = (index, in_dim_rest, concat_dev_num)

576

577 if self.insert_operator(CONCAT_BY_AXIS, args) == Status.FAILED:

578 return False

579

580 del self.map_[index]

581 return True

582

583 def _handle_simple_concat_case(self, index: int, in_dim: Union[int, Tuple[int, ...]],

584 out_dim: Union[int, Tuple[int, ...]]) -> bool:

585 """Handle simple concat case where input dimension is mapped but output is not"""

586 if in_dim == NONE:

587 return False

588

589 if self.out_tensor_map_.get_index_by_value(in_dim) != NONE:

590 return False

591

592 concat_dev_num = self.dev_mat_.get_dim_by_reverse_idx(in_dim)

593 args = (index, in_dim, concat_dev_num)

594

595 if self.insert_operator(CONCAT_BY_AXIS, args) == Status.FAILED:

596 return False

597

598 if out_dim == NONE:

599 del self.map_[index]

600 else:

601 self.map_[index] = NONE

602

603 return True

604

605 def infer_concat_by_axis(self) -> int:

606 """

607 Infers concat operations for the current mapping state.

608

609 Conditions for concat:

610 - Input dimension is mapped but output is unmapped

611 - Device dimension needs consolidation

612

613 Returns:

614 Status.SUCCESS if operations inferred, Status.FAILED on error

615 """

616 keys = list(self.map_.keys())

617 for index in keys:

618 if index not in self.map_:

619 continue

620

621 in_dim = self.map_[index]

622 out_dim = self.out_tensor_map_.get_dim_by_idx(index)

623

624 # Handle tuple concat case

625 if self._handle_tuple_concat_case(index, in_dim, out_dim):

626 continue

627

628 # Handle simple concat case

629 if self._handle_simple_concat_case(index, in_dim, out_dim):

630 continue

631

632 return Status.SUCCESS

633

634 def infer_ops_list(self, rank: int, rank_list: List[int]):

635 """

636 Converts internal operator sequence to executable communication operations.

637

638 Args:

639 rank: Current device rank

640 rank_list: Full list of device ranks in row-major order

641

642 Returns:

643 List of executable communication operations as tuples:

644 - ("all_concat", (dim, size, group))

645 - ("all_split", (dim, size, group))

646 - ("all_to_all", (split_dim, concat_dim, size, group))

647 """

648 if self.infer_redistribution_operator() != Status.SUCCESS:

649 raise RuntimeError("infer redistribution operator failed.")

650

651 ops_list = []

652 for op in self.operator_list_:

653 if op[0] == CONCAT_BY_AXIS:

654 tensor_map = [self.dev_ranks - 1 - d for d in op[1][1]] if isinstance(op[1][1], tuple) \

655 else self.dev_ranks - 1 - op[1][1]

656 group = self.dev_mat_.get_devices_along_dim(rank, rank_list, tensor_map)

657 concat_dim = op[1][0]

658 concat_size = op[1][2]

659 if concat_size == 1:

660 continue

661 ops_list.append(("all_concat", (concat_dim, concat_size, group)))

662 elif op[0] == SPLIT_BY_AXIS:

663 tensor_map = [self.dev_ranks - 1 - d for d in op[1][1]] if isinstance(op[1][1], tuple) \

664 else self.dev_ranks - 1 - op[1][1]

665 group = self.dev_mat_.get_devices_along_dim(rank, rank_list, tensor_map)

666 split_dim = op[1][0]

667 split_size = op[1][2]

668 if split_size == 1:

669 continue

670 ops_list.append(("all_split", (split_dim, split_size, group)))

671 else:

672 tensor_map = [self.dev_ranks - 1 - d for d in op[1][3]] if isinstance(op[1][3], tuple) \

673 else self.dev_ranks - 1 - op[1][3]

674 group = self.dev_mat_.get_devices_along_dim(rank, rank_list, tensor_map)

675 concat_dim = op[1][2]

676 split_dim = op[1][1]

677 permute_size = op[1][0]

678 if permute_size == 1:

679 continue

680 ops_list.append(("all_to_all", (split_dim, concat_dim, permute_size, group)))

681 return ops_list

Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / dtensor / redistribute_infer.py: 84%

345 statements