experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

[XLA] Update layout tiling in some algebraic simplifier cases.

Browse Source 
PiperOrigin-RevId: 315421103
Change-Id: Ic4dee365811f0faa5678d846a1e1412663610062
This commit is contained in:
A. Unique TensorFlower 2020-06-08 22:26:52 -07:00 committed by TensorFlower Gardener
parent cf1b6b3dfe
commit 6ddc7f8d99
1 changed files with 7 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
tensorflow/compiler/xla/service/algebraic_simplifier.cc
View File

@ -2973,6 +2973,7 @@ Status AlgebraicSimplifierVisitor::HandlePad(HloInstruction* pad) {
// slice instruction should all have the same layout. // slice instruction should all have the same layout.
TF_RETURN_IF_ERROR(LayoutUtil::CopyLayoutBetweenShapes( TF_RETURN_IF_ERROR(LayoutUtil::CopyLayoutBetweenShapes(
pad->shape(), nonzero_pad->mutable_shape())); pad->shape(), nonzero_pad->mutable_shape()));
simplifier_->UpdateLayout(nonzero_pad->mutable_shape());
// Second, construct the slice instruction to perform the negative padding. // Second, construct the slice instruction to perform the negative padding.
std::vector<int64> start_indices; std::vector<int64> start_indices;
@ -2999,9 +3000,14 @@ Status AlgebraicSimplifierVisitor::HandlePad(HloInstruction* pad) {
MakeSliceHlo(nonzero_pad, start_indices, end_indices, strides)); MakeSliceHlo(nonzero_pad, start_indices, end_indices, strides));
TF_RETURN_IF_ERROR(LayoutUtil::CopyLayoutBetweenShapes( TF_RETURN_IF_ERROR(LayoutUtil::CopyLayoutBetweenShapes(
pad->shape(), slice->mutable_shape())); pad->shape(), slice->mutable_shape()));
simplifier_->UpdateLayout(slice->mutable_shape());
// Verify that the slice shape matches the pad shape. // Verify that the slice shape matches the pad shape.
TF_RET_CHECK(ShapeUtil::Equal(slice->shape(), pad->shape())); auto equal = Shape::Equal();
if (!options_.is_layout_sensitive()) {
equal.IgnoreTilesInLayout();
}
TF_RET_CHECK(equal(slice->shape(), pad->shape()));
return ReplaceInstruction(pad, slice); return ReplaceInstruction(pad, slice);
} }