Sharding custom calls cannot be hoisted individually in WhileLoopInvariantCodeMotion.

[copybara-service]

Sharding custom calls cannot be hoisted individually in WhileLoopInvariantCodeMotion.

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.