[XLA] Fix #17090 a problem in IrArray::Index::SourceIndexOfTranspose.
Agebraic simplification transforms bitcast-equivalent transpose/reshape instructions to bitcast instructions before IR emission. As such, we should skip the checking on whether a transpose/reshape instruction is bitcast-equivalent or not during IR emission. Remove the call from IrArray::Index::SourceIndexOfTranspose to ShapeUtil::TransposeIsBitcast. Also remove the call from IrArray::Index::SourceIndexOfReshape to ShapeUtil::ReshapeIsBitcast. Remove the calls to ShapeUtil::TransposeIsBitcast and ShapeUtil::ReshapeIsBitcast from NotWorthHoistingIndividually because layout assignment hasn't been done there yet. Instead, returns true when the input is a transpose or reshape instruction, to prevent it from being hoisted out of loops. Add a check to ShapeUtil::TransposeIsBitcast and ShapeUtil::ReshapeIsBitcast to make sure that both input shape and output shape have layouts. Add two test cases. PiperOrigin-RevId: 187093399
Loading
Please sign in to comment