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

Skip ScopedAllocatorOptimizer when the op inputs are _Arg nodes.

Browse Source 
The optimizer adds control inputs from the `_ScopedAllocator` node to the op
input.  But that would cause a dependency from the function body to function
inputs.

PiperOrigin-RevId: 295227438
Change-Id: I9955aef9b22d6b32f7cfb0fe4b97199b502bf131
This commit is contained in:
Ayush Dubey 2020-02-14 14:28:02 -08:00 committed by TensorFlower Gardener
parent c47458b7ba
commit 1edb79fb35
1 changed files with 8 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
tensorflow/core/grappler/optimizers/scoped_allocator_optimizer.cc
View File

@ -518,6 +518,11 @@ class UnaryElementwiseRewriter : public ScopedAllocatorOptimizer::Rewriter {
// input nodes and mark them for allocation from backing tensor.
for (int i = 0; i < inputs.size(); ++i) {
auto& nd = inputs[i];
if (IsArg(*nd.from_node_def)) {
return errors::Internal(
"ScopedAllocatorOptimizer does not work well when the op inputs "
"are _Arg ops; skipping this optimizer for this function");
}
VLOG(2) << "To input " << i << ": " << nd.from_node_def->name()
<< " add control input "
<< "^" << sa_name;
@ -530,8 +535,9 @@ class UnaryElementwiseRewriter : public ScopedAllocatorOptimizer::Rewriter {
node_map->AddOutput(sa_name, nd.from_node_def->name());
}
// We add control edges in order to delay execution of the ScopedAllocatorOp
// until just before first use in order to conserve memory.
// We add control edges in order to (1) delay execution of the
// ScopedAllocatorOp until just before first use in order to conserve memory
// (2) ensure correctness in the presence of control flow related ops.
bool added_delay_edge = false;
for (auto& nd : inputs) {
std::vector<InputDesc> inputs_to_first;