Commit 5b16e707 authored by Tong Shen's avatar Tong Shen Committed by TensorFlower Gardener
Browse files

Fix control input handling for XLA/TPU.

Issue with host loop + TPU rewrite + device loop, before cl/219724472:
a) For device loop, tf.while_loop python code selects a certain Merge node as "pivot", and will add control edges between the "pivot" and nodes in loop cond ("Less" node).
b) AddOp("Less" node) in device loop context ---> AddInnerOp("Less" node) in TPU rewrite context ---> AddOp("Less" node) in TPU rewrite context ---> _RemoveExternalControlEdges("Less" node). The _RemoveExternalControlEdges() function checks "Merge" node's control flow context (which is the device loop context), and "Less" node's control flow context (which is TPU rewrite context at the moment), so it decides "Merge" node is an external control input for "Less" node. It adds an Identity node for the "Merge", and adds control edge between the new Identity node and "Less" node. We use control flow context escaping on the new Identity node, so it does not have any control flow context other than TPU rewrite context, thus is *not* visible to device loop context.
c) At TF execution time, TF complains that the new Identity node has control edge to "Less" node, but they are not in the same while loop context: Identity node does not have any loop context, and "Less" node has device loop context.

How cl/219724472 fixed the issue: in step b), remove the control flow context escaping, so the device loop context sees the new Identity node. Now both Identity node and "Less" node belong to device loop context. But it causes another issue: if there is an control edge between host loop node and device loop node, then we get into an infinite loop: TPU rewrite context adds an Identity -> AddOp() in device loop context -> AddInnerOp() in TPU rewrite context -> AddOp() in TPU rewrite context -> TPU rewrite contet adds an Identity.

How this CL fixed the issue:
1) rollback cl/219724472;
2) in step b), do not let TPU rewrite context think Merge node is an external control input. As a result, the Identity node won't be created, and TF won't complain about the control edge between Identity node and "Less" node.

PiperOrigin-RevId: 231427455
parent f0c49a7f
Loading
Loading
Loading
Loading
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please to comment