From b75ea8de711a4b6fb55124228a563678726e5a9e Mon Sep 17 00:00:00 2001
From: Andy Ly <lyandy@google.com>
Date: Wed, 3 Jun 2020 09:06:05 -0700
Subject: [PATCH] Fix handling of single op islands when breaking up islands.

It is possible for a single op island return different results from the single op (duplicate results, aliased results). This updates the single op island check with WrapsSingleOp instead of checking if the island body only has a single op excluding the terminator.

PiperOrigin-RevId: 314541656
Change-Id: I10bc8e6348a72764579c0ee0f977d100bf069aa8
---
 .../tensorflow/tests/breakup-islands.mlir     | 28 +++++++++++++++++++
 .../tensorflow/translate/breakup-islands.cc   |  2 +-
 2 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/tensorflow/compiler/mlir/tensorflow/tests/breakup-islands.mlir b/tensorflow/compiler/mlir/tensorflow/tests/breakup-islands.mlir
index 61e0772726c..290da5a6135 100644
--- a/tensorflow/compiler/mlir/tensorflow/tests/breakup-islands.mlir
+++ b/tensorflow/compiler/mlir/tensorflow/tests/breakup-islands.mlir
@@ -344,3 +344,31 @@ func @switchn_control_input(%arg1: tensor<i32>) {
   }
   return
 }
+
+// CHECK-LABEL: func @single_op_island_forward_block_arg
+// CHECK: %[[CONST:.*]], %{{.*}} = tf_executor.island wraps "tf.Const"
+// CHECK: tf_executor.fetch %[[CONST]], %arg0
+func @single_op_island_forward_block_arg(%arg0: tensor<?x?x?x?xbf16>) -> (tensor<2048xf32>, tensor<?x?x?x?xbf16>) {
+  %0:2 = tf_executor.graph {
+    %outputs:2, %control = tf_executor.island {
+      %1 = "tf.Const"() {value = dense<0.000000e+00> : tensor<2048xf32>} : () -> tensor<2048xf32>
+      tf_executor.yield %1, %arg0 : tensor<2048xf32>, tensor<?x?x?x?xbf16>
+    }
+    tf_executor.fetch %outputs#0, %outputs#1 : tensor<2048xf32>, tensor<?x?x?x?xbf16>
+  }
+  return %0#0, %0#1 : tensor<2048xf32>, tensor<?x?x?x?xbf16>
+}
+
+// CHECK-LABEL: func @single_op_island_duplicate_result
+// CHECK: %[[CONST:.*]], %{{.*}} = tf_executor.island wraps "tf.Const"
+// CHECK: tf_executor.fetch %[[CONST]], %[[CONST]]
+func @single_op_island_duplicate_result() -> (tensor<2048xf32>, tensor<2048xf32>) {
+  %0:2 = tf_executor.graph {
+    %outputs:2, %control = tf_executor.island {
+      %1 = "tf.Const"() {value = dense<0.000000e+00> : tensor<2048xf32>} : () -> tensor<2048xf32>
+      tf_executor.yield %1, %1 : tensor<2048xf32>, tensor<2048xf32>
+    }
+    tf_executor.fetch %outputs#0, %outputs#1 : tensor<2048xf32>, tensor<2048xf32>
+  }
+  return %0#0, %0#1 : tensor<2048xf32>, tensor<2048xf32>
+}
diff --git a/tensorflow/compiler/mlir/tensorflow/translate/breakup-islands.cc b/tensorflow/compiler/mlir/tensorflow/translate/breakup-islands.cc
index 3245e3b9e6a..7284626c46a 100644
--- a/tensorflow/compiler/mlir/tensorflow/translate/breakup-islands.cc
+++ b/tensorflow/compiler/mlir/tensorflow/translate/breakup-islands.cc
@@ -219,7 +219,7 @@ void BreakUpIslands::BreakUpIsland(
   }
 
   // Skip islands that are already only a single op.
-  if (hasSingleElement(island_body)) return;
+  if (island_op.WrapsSingleOp()) return;
 
   auto control_type = tf_executor::ControlType::get(&getContext());
   auto island_control_inputs = llvm::to_vector<4>(island_op.controlInputs());