Commits on Sep 23, 2024

1. [XLA] Introduce infeed token propagation …

   During computation inlining, specifically loop unrolling, it is posibble for infeeds (and outfeeds) to get reordered in a way that breaks the original scheduling constraints set by the computation boundaries. This is a result of Tensorflow not exposing tokens for these ops to the user, so the input and output tokens end up dangling.
   
   Loop unrolling in XLA can be thought of applying the same function repeatedly to itself, e.g. transforming f(x) into f(f(x)). By pushing the tokens outside the loop body, we can guarantee that the output token of the first infeed will become the input token of the next infeed, thus creating a data dependency chain and preserving the original ordering.
   
   PiperOrigin-RevId: 677572109

   

   vsytch authored and Google-ML-Automation committed Sep 23, 2024
   Configuration menu

   • View commit details
   • Copy full SHA for ff8d1d8
   • Browse repository at this point

   Copy the full SHA

   ff8d1d8 View commit details

   Browse the repository at this point in the history