From b7fb764f9fd8e84ef089d3f3e5a805a01cddde62 Mon Sep 17 00:00:00 2001
From: Smit Hinsu <hinsu@google.com>
Date: Mon, 20 Jul 2020 14:15:14 -0700
Subject: [PATCH] Verify IslandOp body has no block arguments

Also, removed explicit check on empty body which is already verified by trait
SingleBlockImplicitTerminator.

PiperOrigin-RevId: 322222883
Change-Id: Id9344905328a5680a02b8c95fe2805e230780e77
---
 .../compiler/mlir/tensorflow/ir/tf_executor.cc     |  4 ++--
 .../tensorflow/tests/tf_executor_ops_invalid.mlir  | 14 ++++++++++++++
 2 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_executor.cc b/tensorflow/compiler/mlir/tensorflow/ir/tf_executor.cc
index 1b1d5ba6f3b..8db06e83527 100644
--- a/tensorflow/compiler/mlir/tensorflow/ir/tf_executor.cc
+++ b/tensorflow/compiler/mlir/tensorflow/ir/tf_executor.cc
@@ -301,8 +301,8 @@ bool IslandOp::WrapsSingleOp() {
 namespace {
 
 LogicalResult Verify(IslandOp island) {
-  if (island.GetBody().empty())
-    return island.emitOpError() << "expects a non-empty body";
+  if (!island.GetBody().args_empty())
+    return island.emitOpError() << "expects body without any arguments";
 
   Operation &yield = island.GetBody().back();
   if (!isa<YieldOp>(yield))
diff --git a/tensorflow/compiler/mlir/tensorflow/tests/tf_executor_ops_invalid.mlir b/tensorflow/compiler/mlir/tensorflow/tests/tf_executor_ops_invalid.mlir
index 2f034f1bfae..0e9814de137 100644
--- a/tensorflow/compiler/mlir/tensorflow/tests/tf_executor_ops_invalid.mlir
+++ b/tensorflow/compiler/mlir/tensorflow/tests/tf_executor_ops_invalid.mlir
@@ -232,6 +232,20 @@ func @invalid_island(%arg0: tensor<*xf32>, %ctl: !tf_executor.control) {
 
 // -----
 
+// Check that an island body doesn't have any block arguments.
+func @invalid_island(%arg0: tensor<*xf32>, %ctl: !tf_executor.control) {
+  tf_executor.graph {
+    "tf_executor.island"() ({
+      // expected-error@-1 {{expects body without any arguments}}
+      ^entry(%arg: tensor<2xi32>):
+        tf_executor.yield
+    }) : () -> (!tf_executor.control)
+  }
+  return
+}
+
+// -----
+
 // Check that an island body can't be empty.
 func @invalid_island(%arg0: tensor<*xf32>, %ctl: !tf_executor.control) {
   tf_executor.graph {