Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
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
- Loading branch information
1 parent
839aa17
commit 5c90e79
Showing
2 changed files
with
50 additions
and
4 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters