commit e2ddc32097da1d6aefe5e9122572febe5727b7c0
author Berkin Ilbeyi <berkin@google.com> Mon Mar 28 19:29:29 2022 -0700
committer TensorFlower Gardener <gardener@tensorflow.org> Mon Mar 28 19:32:40 2022 -0700
tree 4d6db175ba1be1b66142cfbfea2f548abf1928ea
parent d1241f313c6fcac580194d6805461693c12c51b5

[XLA] Fix copy insertion bug where copies may be inserted at wrong while level in nested whiles

Consider nested while loops like the following:

cond.inner {
  ROOT param.cond.inner = pred[] parameter(0)
}

body.inner {
  param.body.inner = pred[] parameter(0)
  ROOT not = pred[] not(param.body.inner)
}

cond.outer {
  ROOT param.cond.outer = pred[] parameter(0)
}

body.outer {
  ROOT param.cond.outer = pred[] parameter(0)
  while = pred[] while(param.cond.outer), condition=cond.inner, body=body.inner
  after-all = token[] after-all()
  outfeed = token[] outfeed(while, after-all)
}

ENTRY TestComputation {
  entry_param = pred[] parameter(0)
  while = pred[] while(entry_param), condition=cond.outer, body=body.outer
  ROOT not = pred[] not(while)
}

Because the inner while loop modifies the value, a copy needs to be inserted.
Due to the SSA semantics of pre-copy-inserted HLO, this copy should have been
inserted in the outer while loop's body, to make sure the subsequent iteration
of the outer while loop gets a fresh copy of the entry parameter. However, the
copy insertion pass was inserting the copy in the entry computation. Fixed this
bug by making sure if a value is the root of a computation a second value is
defined in, the first value's live range is not before the second value, so this
copy in the outer while body cannot be removed.

Also removed workarounds in tuple analysis (cl/415649802) and lingvo gshard
(cl/409279977) since these are no longer needed.

PiperOrigin-RevId: 437913642