From 763ff559702ea6cf21551a89dcbbaa3fefd4c0df Mon Sep 17 00:00:00 2001
From: Geoffrey Martin-Noble <gcmn@google.com>
Date: Thu, 1 Apr 2021 13:03:17 -0700
Subject: [PATCH] Restore SingleBlockImplicitTerminator verification to
 mhlo.while

The internal users have been cleaned up, so we can roll this forward again.

PiperOrigin-RevId: 366313960
---
 include/mlir-hlo/Dialect/mhlo/IR/hlo_ops.td | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/include/mlir-hlo/Dialect/mhlo/IR/hlo_ops.td b/include/mlir-hlo/Dialect/mhlo/IR/hlo_ops.td
index 8b338e4..b3eabbf 100644
--- a/include/mlir-hlo/Dialect/mhlo/IR/hlo_ops.td
+++ b/include/mlir-hlo/Dialect/mhlo/IR/hlo_ops.td
@@ -578,7 +578,8 @@ def HLO_CaseOp: HLO_Op<"case", [
 
 def HLO_WhileOp: HLO_Op<"while", [
       RecursiveSideEffects,
-      SameOperandsAndResultType
+      SameOperandsAndResultType,
+      SingleBlockImplicitTerminator<"ReturnOp">
     ]>, BASE_HLO_WhileOp {
   let arguments = (ins HLO_TensorOrTuple:$val);