experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Remove the use of SetDynamicBinding in tf2xla bridge.

Browse Source 
- 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
This commit is contained in:
Yunxing Dai 2020-08-17 11:08:25 -07:00 committed by TensorFlower Gardener
parent b9fce9ec2e
commit 27da5d74dc
2 changed files with 11 additions and 23 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

33
tensorflow/compiler/tf2xla/xla_compiler.cc
View File

@ -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<int>::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);
}
}

1
tensorflow/compiler/xla/client/xla_builder.h
View File

@ -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,