From d0a7bb9b1ad768dfd9b08f66a0e744584b49181f Mon Sep 17 00:00:00 2001
From: Matthias Springer <springerm@google.com>
Date: Sat, 4 Feb 2023 15:44:20 +0100
Subject: [PATCH] [mlir][SCF] Disallow multiple blocks in scf.if "else" region

The "then" region allows only a single block. The same should be the case for the "else" region.

Differential Revision: https://reviews.llvm.org/D143253
---
 mlir/include/mlir/Dialect/SCF/IR/SCFOps.td | 22 ++++++++++++++--------
 mlir/include/mlir/IR/OpBase.td             |  5 +++++
 2 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
index b3b8825..7e2d337 100644
--- a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
+++ b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
@@ -630,8 +630,8 @@ def IfOp : SCF_Op<"if", [DeclareOpInterfaceMethods<RegionBranchOpInterface, [
     }
     ```
 
-    `scf.if` may also return results that are defined in its regions. The
-    values defined are determined by which execution path is taken.
+    `scf.if` may also produce results. Which values are returned depends on
+    which execution path is taken.
 
     Example:
 
@@ -647,11 +647,13 @@ def IfOp : SCF_Op<"if", [DeclareOpInterfaceMethods<RegionBranchOpInterface, [
     }
     ```
 
-    `scf.if` regions are always terminated with "scf.yield". If "scf.if"
-    defines no values, the "scf.yield" can be left out, and will be inserted
-    implicitly. Otherwise, it must be explicit.
-    Also, if "scf.if" defines one or more values, the 'else' block cannot be
-    omitted.
+    The "then" region has exactly 1 block. The "else" region may have 0 or 1
+    block. In case the `scf.if` produces results, the "else" region must also
+    have exactly 1 block.
+
+    The blocks are always terminated with `scf.yield`. If `scf.if` defines no
+    values, the `scf.yield` can be left out, and will be inserted implicitly.
+    Otherwise, it must be explicit.
 
     Example:
 
@@ -660,10 +662,14 @@ def IfOp : SCF_Op<"if", [DeclareOpInterfaceMethods<RegionBranchOpInterface, [
       ...
     }
     ```
+
+    The types of the yielded values must match the result types of the
+    `scf.if`.
   }];
   let arguments = (ins I1:$condition);
   let results = (outs Variadic<AnyType>:$results);
-  let regions = (region SizedRegion<1>:$thenRegion, AnyRegion:$elseRegion);
+  let regions = (region SizedRegion<1>:$thenRegion,
+                        MaxSizedRegion<1>:$elseRegion);
 
   let skipDefaultBuilders = 1;
   let builders = [
diff --git a/mlir/include/mlir/IR/OpBase.td b/mlir/include/mlir/IR/OpBase.td
index d307beb..7ef6268 100644
--- a/mlir/include/mlir/IR/OpBase.td
+++ b/mlir/include/mlir/IR/OpBase.td
@@ -1762,6 +1762,11 @@ class MinSizedRegion<int numBlocks> : Region<
   CPred<"::llvm::hasNItemsOrMore($_self, " # numBlocks # ")">,
   "region with at least " # numBlocks # " blocks">;
 
+// A region with at most the given number of blocks.
+class MaxSizedRegion<int numBlocks> : Region<
+  CPred<"::llvm::hasNItemsOrLess($_self, " # numBlocks # ")">,
+  "region with at most " # numBlocks # " blocks">;
+
 // A variadic region constraint. It expands to zero or more of the base region.
 class VariadicRegion<Region region>
   : Region<region.predicate, region.summary>;
-- 
2.7.4