[XLA] Fix #17090 a problem in IrArray::Index::SourceIndexOfTranspose.
authorBixia Zheng <bixia@google.com>
Mon, 26 Feb 2018 23:42:52 +0000 (15:42 -0800)
committerGunhan Gulsoy <gunan@google.com>
Tue, 27 Feb 2018 22:33:33 +0000 (14:33 -0800)
commit260f5b8fe144cd369fde755739806449a2901252
treec39ec846e0c2fc6f0ba55216a31924676cec9229
parentacf78b20f71dd8c3a928b1f12ea4de6f5028fc48
[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
tensorflow/compiler/xla/service/layout_assignment_test.cc
tensorflow/compiler/xla/service/llvm_ir/ir_array.cc
tensorflow/compiler/xla/service/while_loop_invariant_code_motion.cc
tensorflow/compiler/xla/shape_util.cc
tensorflow/compiler/xla/shape_util.h