Commits on May 1, 2024

1. Sharding custom calls cannot be hoisted individually in `WhileLoopInv… …

   …ariantCodeMotion`.
   
   A sharding custom call annotates sharding to its operand. If we move a single sharding custom call out of the while body, the sharding annotation may not take effect as expected.
   
   Taking the follow while body as an example,
   ```
   body {
     p_body = (f32[2], f32[2], s32[]) parameter(0)
     gte.0 = f32[2] get-tuple-element(p_body), index=0
     gte.1 = f32[2] get-tuple-element(p_body), index=1
     sharding.0 = f32[2] custom-call(gte.0), custom_call_target="Sharding", sharding={devices=[2]<=[2]}
     sharding.1 = f32[2] custom-call(gte.1), custom_call_target="Sharding", sharding={replicated}
     add.0 = f32[2] add(sharding.0, sharding.1)
     gte.2 = s32[] get-tuple-element(p_body), index=2
     const = s32[] constant(1)
     add.1 = s32[] add(gte.2, const)
     ROOT root = (f32[2], f32[2], s32[]) tuple(gte.0, add.0, add.1)
   }
   ```
   
   Before this cl, `WhileLoopInvariantCodeMotion` moves `sharding.0` out of the body and keeps `sharding.1` in the body. With this cl, `WhileLoopInvariantCodeMotion` does not modify this body.
   
   PiperOrigin-RevId: 629637435

   

   ZixuanJiang authored and tensorflower-gardener committed May 1, 2024
   Configuration menu

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

   Copy the full SHA

   5c90e79 View commit details

   Browse the repository at this point in the history