| 30 | namespace xla { |
| 31 | |
| 32 | StatusOr<bool> DynamicIndexSplitter::Run(HloModule* module) { |
| 33 | bool changed = false; |
| 34 | |
| 35 | std::vector<HloComputation*> computations = |
| 36 | module->MakeNonfusionComputations(); |
| 37 | for (HloComputation* computation : computations) { |
| 38 | for (HloInstruction* dynamic_op : computation->MakeInstructionPostOrder()) { |
| 39 | switch (dynamic_op->opcode()) { |
| 40 | case HloOpcode::kDynamicSlice: |
| 41 | case HloOpcode::kDynamicUpdateSlice: |
| 42 | break; |
| 43 | default: |
| 44 | continue; |
| 45 | } |
| 46 | auto parent = dynamic_op->parent(); |
| 47 | bool is_update = dynamic_op->opcode() == HloOpcode::kDynamicUpdateSlice; |
| 48 | int64 num_indices = dynamic_op->operand(0)->shape().rank(); |
| 49 | |
| 50 | if (num_indices == 0) { |
| 51 | // If the operand rank is 0, directly replace R0 DS/DUS with the |
| 52 | // operand (for DS) or update (for DUS). |
| 53 | if (is_update) { |
| 54 | TF_CHECK_OK(parent->ReplaceInstruction( |
| 55 | dynamic_op, dynamic_op->mutable_operand(1))); |
| 56 | } else { |
| 57 | TF_CHECK_OK(parent->ReplaceInstruction( |
| 58 | dynamic_op, dynamic_op->mutable_operand(0))); |
| 59 | } |
| 60 | changed = true; |
| 61 | continue; |
| 62 | } |
| 63 | |
| 64 | int64 index_operand_number = Cast<HloDynamicIndexInstruction>(dynamic_op) |
| 65 | ->first_index_operand_number(); |
| 66 | auto index_operand = dynamic_op->mutable_operand(index_operand_number); |
| 67 | if (ShapeUtil::IsScalar(index_operand->shape())) { |
| 68 | // This DS/DUS already uses scalar indices. |
| 69 | continue; |
| 70 | } |
| 71 | TF_RET_CHECK(index_operand->shape().rank() == 1); |
| 72 | auto index_element_type = index_operand->shape().element_type(); |
| 73 | std::vector<HloInstruction*> index_array; |
| 74 | for (int64 dim = 0; dim < num_indices; ++dim) { |
| 75 | auto slice = parent->AddInstruction(HloInstruction::CreateSlice( |
| 76 | ShapeUtil::MakeShape(index_element_type, {1}), index_operand, {dim}, |
| 77 | {dim + 1}, {1})); |
| 78 | auto bitcast = parent->AddInstruction(HloInstruction::CreateReshape( |
| 79 | ShapeUtil::MakeShape(index_element_type, {}), slice)); |
| 80 | index_array.push_back(bitcast); |
| 81 | } |
| 82 | auto new_dynamic_op = |
| 83 | is_update |
| 84 | ? HloInstruction::CreateDynamicUpdateSlice( |
| 85 | dynamic_op->shape(), dynamic_op->mutable_operand(0), |
| 86 | dynamic_op->mutable_operand(1), absl::MakeSpan(index_array)) |
| 87 | : HloInstruction::CreateDynamicSlice( |
| 88 | dynamic_op->shape(), dynamic_op->mutable_operand(0), |
| 89 | absl::MakeSpan(index_array), |
nothing calls this directly
no test coverage detected