From 3b94e9cfdb33c58b215b8157c6ef34af4fb4fdfe Mon Sep 17 00:00:00 2001
From: Yunxing Dai <yunxing@google.com>
Date: Sat, 16 Jan 2021 10:15:40 -0800
Subject: [PATCH] Support 0-sized shape with dynamic reshapes.

- Return finer granule control in CommonFactors when input and output are the same and one of dimensions is zero.
- Skip instruction with dynamic size in ZeroSizedHloElimination.

PiperOrigin-RevId: 352189050
Change-Id: Ieda2d8c89d20e66670b3cd85484f4e1e97bc4c93
---
 .../compiler/xla/service/zero_sized_hlo_elimination.cc   | 3 ++-
 tensorflow/compiler/xla/util.cc                          | 9 ++++++++-
 tensorflow/compiler/xla/util.h                           | 6 ++++--
 tensorflow/compiler/xla/util_test.cc                     | 1 +
 4 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/tensorflow/compiler/xla/service/zero_sized_hlo_elimination.cc b/tensorflow/compiler/xla/service/zero_sized_hlo_elimination.cc
index 4c221e2c116..975f687fac4 100644
--- a/tensorflow/compiler/xla/service/zero_sized_hlo_elimination.cc
+++ b/tensorflow/compiler/xla/service/zero_sized_hlo_elimination.cc
@@ -36,7 +36,8 @@ StatusOr<bool> ZeroSizedHloElimination::Run(HloModule* module) {
         continue;
       }
       if (comp->IsSafelyRemovable(instruction) &&
-          ShapeUtil::IsZeroElementArray(instruction->shape())) {
+          ShapeUtil::IsZeroElementArray(instruction->shape()) &&
+          instruction->shape().is_static()) {
         // If the instruction doesn't have a layout, use a default layout for
         // the literal.
         Shape shape = instruction->shape();
diff --git a/tensorflow/compiler/xla/util.cc b/tensorflow/compiler/xla/util.cc
index f39bd269ef4..fdcb33f784b 100644
--- a/tensorflow/compiler/xla/util.cc
+++ b/tensorflow/compiler/xla/util.cc
@@ -265,11 +265,18 @@ int64 Product(absl::Span<const int64> xs) {
 absl::InlinedVector<std::pair<int64, int64>, 8> CommonFactors(
     absl::Span<const int64> a, absl::Span<const int64> b) {
   CHECK_EQ(Product(a), Product(b));
+  absl::InlinedVector<std::pair<int64, int64>, 8> bounds;
+  if (absl::c_equal(a, b)) {
+    bounds.reserve(a.size() + 1);
+    for (int64 i = 0; i <= a.size(); ++i) {
+      bounds.emplace_back(i, i);
+    }
+    return bounds;
+  }
   if (0 == Product(a)) {
     return {std::make_pair(0, 0), std::make_pair(a.size(), b.size())};
   }
 
-  absl::InlinedVector<std::pair<int64, int64>, 8> bounds;
   for (int64 i = 0, j = 0, prior_i = -1, prior_j = -1, partial_size_a = 1,
              partial_size_b = 1;
        ;) {
diff --git a/tensorflow/compiler/xla/util.h b/tensorflow/compiler/xla/util.h
index 44a5bf4ea33..a3fbbd51805 100644
--- a/tensorflow/compiler/xla/util.h
+++ b/tensorflow/compiler/xla/util.h
@@ -502,8 +502,10 @@ int64 Product(absl::Span<const int64> xs);
 //         b[j_k] × b[j_k + 1] × ... × b[j_(k+1) - 1]
 // where `CommonFactors(a, b)[CommonFactors(a, b).size - 1] = (a.size, b.size)`
 //
-// If the given shapes have non-zero size, returns the bounds of the shortest
-// possible such subsequences; else, returns `{(0, 0), (a.size, b.size)}`.
+// If input and output are the same, return {(0, 0), {1, 1}, ... {a.size,
+// b.size}}, otherwise if the given shapes have non-zero size, returns the
+// bounds of the shortest possible such subsequences; else, returns `{(0, 0),
+// (a.size, b.size)}`.
 absl::InlinedVector<std::pair<int64, int64>, 8> CommonFactors(
     absl::Span<const int64> a, absl::Span<const int64> b);
 
diff --git a/tensorflow/compiler/xla/util_test.cc b/tensorflow/compiler/xla/util_test.cc
index 6f60e241d92..17d6400c75f 100644
--- a/tensorflow/compiler/xla/util_test.cc
+++ b/tensorflow/compiler/xla/util_test.cc
@@ -74,6 +74,7 @@ TEST(UtilTest, CommonFactors) {
     absl::InlinedVector<std::pair<int64, int64>, 8> expected;
   } test_cases[] = {
       {/*.a =*/{0}, /*.b =*/{0}, /*.expected =*/{{0, 0}, {1, 1}}},
+      {/*.a =*/{0, 1}, /*.b =*/{0, 1}, /*.expected =*/{{0, 0}, {1, 1}, {2, 2}}},
       {/*.a =*/{}, /*.b =*/{}, /*.expected =*/{{0, 0}}},
       {/*.a =*/{2, 5, 1, 3},
        /*.b =*/{1, 10, 3, 1},