commit 796af3d216465c9545062fcf5cdd443af6dd08af
author A. Unique TensorFlower <gardener@tensorflow.org> Sat Feb 13 10:01:26 2021 -0800
committer TensorFlower Gardener <gardener@tensorflow.org> Sat Feb 13 10:07:51 2021 -0800
tree 49377cc1c22ce9383c46c60da74bd0e3c6402cbb
parent 37406182181eaf0356fa832b26a3a20e8f2f78a1

We just need to copy the tiling info of result shape in layout assignment instead
the whole shape since the shapes already match.

PiperOrigin-RevId: 357367492
Change-Id: I05be8ebe78265f6d5d4de00c093b8008631275b2

tensorflow/compiler/xla/service/layout_assignment.cc

diff --git a/tensorflow/compiler/xla/service/layout_assignment.cc b/tensorflow/compiler/xla/service/layout_assignment.cc
index 4882c5d..72561cb 100644
--- a/tensorflow/compiler/xla/service/layout_assignment.cc
+++ b/tensorflow/compiler/xla/service/layout_assignment.cc
@@ -1949,9 +1949,19 @@
                                   computation->root_instruction()));
       computation->set_root_instruction(new_root);
     } else {
-      // Use the specified shape including tiling info in layout.
-      *(computation->root_instruction()->mutable_shape()) =
-          constraints.ResultLayout()->shape();
+      // Copy the specified tiling info.
+      auto assign_tiling = [&constraints](xla::Shape* subshape,
+                                          const xla::ShapeIndex& index) {
+        if (subshape->IsArray()) {
+          const Shape& result_shape = ShapeUtil::GetSubshape(
+              constraints.ResultLayout()->shape(), index);
+          subshape->mutable_layout()->mutable_tiles()->assign(
+              result_shape.layout().tiles().begin(),
+              result_shape.layout().tiles().end());
+        }
+      };
+      xla::ShapeUtil::ForEachMutableSubshape(
+          computation->root_instruction()->mutable_shape(), assign_tiling);
     }
   }
   return Status::OK();