| 970 | } // namespace |
| 971 | |
| 972 | StatusOr<bool> DynamicPadder::Run(HloModule* module) { |
| 973 | bool changed = false; |
| 974 | VLOG(2) << "Pre DynamicPadder HLO:"; |
| 975 | |
| 976 | // Removes dynamic dimensions on parameters if there is already a binding for |
| 977 | // it. We do this because we have two different APIs to express a dynamic |
| 978 | // dimension: |
| 979 | // |
| 980 | // 1. Dynamic dimension as specificed directly in the shape -- Needed for |
| 981 | // Pytorch. |
| 982 | // |
| 983 | // 2. Dynamic dimension using dynamic parameter binding object. This |
| 984 | // is needed for tensorflow. |
| 985 | // |
| 986 | // For case 1, we will insert "pad-to-static" instruction in the |
| 987 | // beginning of xla execution, to make it into a static layout. |
| 988 | // |
| 989 | // For case 2, since it already has a static layout, we remove the |
| 990 | // dynamic dimension. |
| 991 | // |
| 992 | // TODO(b/145140571): Convert all API invocations to case 1. |
| 993 | // |
| 994 | TF_RETURN_IF_ERROR(module->dynamic_parameter_binding().ForEachBinding( |
| 995 | [&](const DynamicParameterBinding::DynamicParameter& dynamic_parameter, |
| 996 | const DynamicParameterBinding::DynamicDimension& dynamic_dimension) |
| 997 | -> Status { |
| 998 | HloInstruction* parameter = |
| 999 | module->entry_computation()->parameter_instruction( |
| 1000 | dynamic_dimension.parameter_num); |
| 1001 | ShapeUtil::UpdateDynamicDimension(parameter->mutable_shape(), |
| 1002 | dynamic_dimension.parameter_index, |
| 1003 | dynamic_dimension.dimension, false); |
| 1004 | return Status::OK(); |
| 1005 | })); |
| 1006 | |
| 1007 | TF_RETURN_IF_ERROR(InsertPadToStaticAfterModuleInputs(module)); |
| 1008 | TF_ASSIGN_OR_RETURN(DynamicDimensionInference dynamic_dimension_inference, |
| 1009 | DynamicDimensionInference::Run(module)); |
| 1010 | |
| 1011 | for (HloComputation* computation : module->computations()) { |
| 1012 | for (HloInstruction* inst : computation->MakeInstructionPostOrder()) { |
| 1013 | if (inst->opcode() == HloOpcode::kConcatenate) { |
| 1014 | TF_ASSIGN_OR_RETURN( |
| 1015 | changed, RewriteDynamicConcat(inst, &dynamic_dimension_inference)); |
| 1016 | continue; |
| 1017 | } |
| 1018 | if (inst->opcode() == HloOpcode::kSort) { |
| 1019 | TF_ASSIGN_OR_RETURN( |
| 1020 | changed, RewriteDynamicSort(inst, &dynamic_dimension_inference)); |
| 1021 | continue; |
| 1022 | } |
| 1023 | for (int64 operand_num = 0; operand_num < inst->operand_count(); |
| 1024 | ++operand_num) { |
| 1025 | HloInstruction* original_operand = inst->mutable_operand(operand_num); |
| 1026 | HloInstruction* operand = original_operand; |
| 1027 | if (!operand->shape().IsArray()) { |
| 1028 | continue; |
| 1029 | } |
nothing calls this directly
no test coverage detected