[XLA] Avoid implicit layout change for non-layout-changing instructions.
Add routine InstructionRequiresInputLayoutEqualToOutputLayout to layout assignment and use it to decide whether the output layout of an instruction need to be propagated to the operands of the instruction during layout assignment. This also fixes the performance issue exposed in b/112646847 where a kSlice instruction with implicit layout change is fused to a multiple output fusion for reduction and causes memory throttle. Add test cases. PiperOrigin-RevId: 210622658
Loading
Please sign in to comment