From 27da5d74dc25ab19db692df00606b90257250f63 Mon Sep 17 00:00:00 2001 From: Yunxing Dai Date: Mon, 17 Aug 2020 11:08:25 -0700 Subject: [PATCH] Remove the use of SetDynamicBinding in tf2xla bridge. - Replace SetDynamicBinding with SetDimensionSize models the information into the IR. Makes problems easier to reproduce by just looking at the HLO graph. - This one of the last few places that use SetDynamicBinding, after the clean up, we should be able to replace this old API. PiperOrigin-RevId: 327057424 Change-Id: I7fbadef18a9cd076c12fc61a53310311498416a0 --- tensorflow/compiler/tf2xla/xla_compiler.cc | 33 ++++++-------------- tensorflow/compiler/xla/client/xla_builder.h | 1 + 2 files changed, 11 insertions(+), 23 deletions(-) diff --git a/tensorflow/compiler/tf2xla/xla_compiler.cc b/tensorflow/compiler/tf2xla/xla_compiler.cc index 635b7170d82..d953739ce0c 100644 --- a/tensorflow/compiler/tf2xla/xla_compiler.cc +++ b/tensorflow/compiler/tf2xla/xla_compiler.cc @@ -990,20 +990,6 @@ Status XlaCompiler::BuildArguments( tuple = xla::Parameter(builder, 0, (*input_shapes)[0], "arg_tuple"); } - for (int i = 0, end = input_to_args->size(); i < end; ++i) { - const XlaCompiler::Argument& arg = args[input_to_args->at(i)]; - for (const auto& dim_and_arg_num : arg.dynamic_dim_to_arg_num_map) { - int dynamic_size_param_index = arg_to_inputs.at(dim_and_arg_num.second); - VLOG(1) << "Setting dynamic binding " << i << " -> " - << dynamic_size_param_index; - - TF_RETURN_IF_ERROR(builder->SetDynamicBinding( - /*dynamic_size_param_num=*/0, {dynamic_size_param_index}, - /*target_param_num=*/0, /*target_param_index=*/{i}, - dim_and_arg_num.first)); - } - } - for (std::vector::size_type i = 0; i < input_to_args->size(); ++i) { auto it = arg_shardings.find(i); xla::XlaScopedShardingAssignment assign_sharding( @@ -1035,16 +1021,17 @@ Status XlaCompiler::BuildArguments( absl::StrCat("arg", i)); } } + } - for (int i = 0, end = input_to_args->size(); i < end; ++i) { - const XlaCompiler::Argument& arg = args[input_to_args->at(i)]; - for (const auto& dim_and_arg_num : arg.dynamic_dim_to_arg_num_map) { - int dynamic_size_param_index = arg_to_inputs.at(dim_and_arg_num.second); - TF_RETURN_IF_ERROR(builder->SetDynamicBinding( - /*dynamic_size_param_num=*/dynamic_size_param_index, {}, - /*target_param_num=*/i, /*target_param_index=*/{}, - dim_and_arg_num.first)); - } + for (int i = 0, end = input_to_args->size(); i < end; ++i) { + const XlaCompiler::Argument& arg = args[input_to_args->at(i)]; + for (const auto& dim_and_arg_num : arg.dynamic_dim_to_arg_num_map) { + int dynamic_size_param_index = arg_to_inputs.at(dim_and_arg_num.second); + VLOG(1) << "Setting dynamic size " << i << " -> " + << dynamic_size_param_index; + arg_handles[i] = xla::SetDimensionSize( + arg_handles[i], arg_handles[dynamic_size_param_index], + dim_and_arg_num.first); } } diff --git a/tensorflow/compiler/xla/client/xla_builder.h b/tensorflow/compiler/xla/client/xla_builder.h index a044ec443ef..1bd613e73dd 100644 --- a/tensorflow/compiler/xla/client/xla_builder.h +++ b/tensorflow/compiler/xla/client/xla_builder.h @@ -366,6 +366,7 @@ class XlaBuilder { // // TODO(b/119520625): Remove this API once we have more dynamic shape infra // ready. + ABSL_DEPRECATED("Use SetDimensionSize to set a dynamic dimension.") Status SetDynamicBinding(int64 dynamic_size_param_num, ShapeIndex dynamic_size_param_index, int64 target_param_num,