| 626 | } |
| 627 | |
| 628 | StatusOr<bool> RewriteDynamicSort( |
| 629 | HloInstruction* hlo, |
| 630 | DynamicDimensionInference* dynamic_dimension_inference) { |
| 631 | HloInstruction* dynamic_size = nullptr; |
| 632 | HloSortInstruction* sort = Cast<HloSortInstruction>(hlo); |
| 633 | HloComputation* comp = hlo->parent(); |
| 634 | int64 sort_dim = sort->sort_dimension(); |
| 635 | // Find the dynamic dimension in the operand. |
| 636 | for (auto* operand : sort->operands()) { |
| 637 | if (dynamic_size == nullptr) { |
| 638 | dynamic_size = |
| 639 | dynamic_dimension_inference->GetDynamicSize(operand, {}, sort_dim); |
| 640 | } |
| 641 | } |
| 642 | |
| 643 | if (dynamic_size == nullptr) { |
| 644 | // Not a dynamic sort, ignore. |
| 645 | return false; |
| 646 | } |
| 647 | |
| 648 | Shape operand_shape = |
| 649 | ShapeUtil::ChangeElementType(sort->operand(0)->shape(), S32); |
| 650 | HloInstruction* iota = |
| 651 | comp->AddInstruction(HloInstruction::CreateIota(operand_shape, sort_dim)); |
| 652 | HloInstruction* dynamic_size_broadcasted = comp->AddInstruction( |
| 653 | HloInstruction::CreateBroadcast(operand_shape, dynamic_size, {})); |
| 654 | HloInstruction* lt = comp->AddInstruction(HloInstruction::CreateCompare( |
| 655 | ShapeUtil::ChangeElementType(operand_shape, PRED), iota, |
| 656 | dynamic_size_broadcasted, ComparisonDirection::kLt)); |
| 657 | sort->AppendOperand(lt); |
| 658 | |
| 659 | const int64 param_number_before_rewritten = |
| 660 | sort->called_computations()[0]->num_parameters(); |
| 661 | auto new_param_0 = HloInstruction::CreateParameter( |
| 662 | param_number_before_rewritten, ShapeUtil::MakeScalarShape(PRED), |
| 663 | "inbound_lhs"); |
| 664 | auto new_param_1 = HloInstruction::CreateParameter( |
| 665 | param_number_before_rewritten + 1, ShapeUtil::MakeScalarShape(PRED), |
| 666 | "inbound_rhs"); |
| 667 | std::vector<const HloInstruction*> extra_parameters{new_param_0.get(), |
| 668 | new_param_1.get()}; |
| 669 | HloComputation* sort_comp = sort->parent()->parent()->AddEmbeddedComputation( |
| 670 | sort->called_computations()[0]->CloneWithReplacements( |
| 671 | /*replacements=*/absl::flat_hash_map< |
| 672 | const HloInstruction*, std::unique_ptr<HloInstruction>>(), |
| 673 | extra_parameters)); |
| 674 | auto inbound_lhs = |
| 675 | sort_comp->parameter_instruction(param_number_before_rewritten); |
| 676 | auto inbound_rhs = |
| 677 | sort_comp->parameter_instruction(param_number_before_rewritten + 1); |
| 678 | sort->ReplaceCalledComputations( |
| 679 | [&](HloComputation* comp) { return sort_comp; }); |
| 680 | |
| 681 | // inbound_lhs & (sort_comp | !in_bound_rhs) |
| 682 | // Select the lhs if it is in bounds and the rhs is out of bounds or the |
| 683 | // sort_comp returns true. |
| 684 | auto out_of_bound_rhs = sort_comp->AddInstruction(HloInstruction::CreateUnary( |
| 685 | ShapeUtil::MakeScalarShape(PRED), HloOpcode::kNot, inbound_rhs)); |
no test coverage detected