PartitionId and ReplicaId While Loop Invariant Code Motion

74 views
Skip to first unread message

Patrick Toulme

unread,
Apr 24, 2024, 6:49:57 PMApr 24
to OpenXLA Discuss
Hello, 
I am attempting to hoist HloOpcode::kPartitionId and HloOpcode::kReplicaId out of while loops. Are there any OpenXLA passes to do this? I was looking at the below pass with specifying the opcodes but have not tested it yet. 


Any thoughts are appreciated. 

Best, 
Patrick 

Patrick Toulme

unread,
Apr 24, 2024, 9:14:53 PMApr 24
to OpenXLA Discuss, Patrick Toulme
Update that pass did not work. My new idea is to write a pass that only lifts partition id and replica id out of while bodies and conditionals and then run that to fixed point. Eventually the partition and replica ids will be in the main computation where I need them to be.
Reply all
Reply to author
Forward
0 new messages