The operands of a call must match the layouts of parameters in the ComputationLayout, and the call instruction itself must match the result layout in the ComputationLayout.
| 676 | // ComputationLayout, and the call instruction itself must match the result |
| 677 | // layout in the ComputationLayout. |
| 678 | Status CheckCallLayout(HloInstruction* call, |
| 679 | const ComputationLayout& computation_layout) { |
| 680 | HloComputation* computation = call->to_apply(); |
| 681 | TF_RET_CHECK(computation->num_parameters() == call->operand_count()); |
| 682 | for (int64 i = 0; i < computation->num_parameters(); ++i) { |
| 683 | TF_RET_CHECK(computation_layout.parameter_layout(i).MatchesLayoutInShape( |
| 684 | call->operand(i)->shape(), /*minor_to_major_only=*/true)); |
| 685 | } |
| 686 | TF_RET_CHECK(computation_layout.result_layout().MatchesLayoutInShape( |
| 687 | call->shape(), /*minor_to_major_only=*/true)); |
| 688 | return Status::OK(); |
| 689 | } |
| 690 | |
| 691 | // Operands of layout-constrained custom calls must match the expected |
| 692 | // constrained layouts. |
no test coverage detected