PartitionId and ReplicaId While Loop Invariant Code Motion

[Patrick Toulme]

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. 

 https://github.com/openxla/xla/blob/4797c08a3fd391caf18a26a82003478274ef41da/xla/service/while_loop_expensive_invariant_code_motion.h

Any thoughts are appreciated. 

Best, 

Patrick 

[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.