[MLIR] Add a NoRegionArguments trait

- This trait will verify that all regions attached to an Op have no arguments
- Fixes https://bugs.llvm.org/show_bug.cgi?id=46521 : Add trait NoRegionArguments

Differential Revision: https://reviews.llvm.org/D83016

mlir/include/mlir/Dialect/Affine/IR/AffineOps.td

@@ -278,7 +278,8 @@ def AffineForOp : Affine_Op<"for",
 }
 
 def AffineIfOp : Affine_Op<"if",
-                           [ImplicitAffineTerminator, RecursiveSideEffects]> {
+                           [ImplicitAffineTerminator, RecursiveSideEffects,
+                            NoRegionArguments]> {
   let summary = "if-then-else operation";
   let description = [{
     Syntax:

mlir/include/mlir/Dialect/SCF/SCFOps.td

@@ -183,7 +183,8 @@ def ForOp : SCF_Op<"for",
 
 def IfOp : SCF_Op<"if",
       [DeclareOpInterfaceMethods<RegionBranchOpInterface>,
-       SingleBlockImplicitTerminator<"YieldOp">, RecursiveSideEffects]> {
+       SingleBlockImplicitTerminator<"YieldOp">, RecursiveSideEffects,
+       NoRegionArguments]> {
   let summary = "if-then-else operation";
   let description = [{
     The `scf.if` operation represents an if-then-else construct for

mlir/include/mlir/IR/Module.h

@@ -33,7 +33,7 @@ class ModuleOp
           OpTrait::IsIsolatedFromAbove, OpTrait::AffineScope,
           OpTrait::SymbolTable,
           OpTrait::SingleBlockImplicitTerminator<ModuleTerminatorOp>::Impl,
-          SymbolOpInterface::Trait> {
+          SymbolOpInterface::Trait, OpTrait::NoRegionArguments> {
 public:
   using Op::Op;
   using Op::print;

mlir/include/mlir/IR/OpBase.td

@@ -1737,6 +1737,9 @@ def AttrSizedOperandSegments : NativeOpTrait<"AttrSizedOperandSegments">;
 // should be named as `result_segment_sizes`.
 def AttrSizedResultSegments  : NativeOpTrait<"AttrSizedResultSegments">;
 
+// Op attached regions have no arguments
+def NoRegionArguments : NativeOpTrait<"NoRegionArguments">;
+
 //===----------------------------------------------------------------------===//
 // OpInterface definitions
 //===----------------------------------------------------------------------===//

mlir/include/mlir/IR/OpDefinition.h

@@ -400,6 +400,7 @@ LogicalResult verifyNSuccessors(Operation *op, unsigned numSuccessors);
 LogicalResult verifyAtLeastNSuccessors(Operation *op, unsigned numSuccessors);
 LogicalResult verifyOperandSizeAttr(Operation *op, StringRef sizeAttrName);
 LogicalResult verifyResultSizeAttr(Operation *op, StringRef sizeAttrName);
+LogicalResult verifyNoRegionArguments(Operation *op);
 } // namespace impl
 
 /// Helper class for implementing traits.  Clients are not expected to interact
@@ -1202,6 +1203,15 @@ public:
   }
 };
 
+/// This trait provides a verifier for ops that are expecting their regions to
+/// not have any arguments
+template <typename ConcrentType>
+struct NoRegionArguments : public TraitBase<ConcrentType, NoRegionArguments> {
+  static LogicalResult verifyTrait(Operation *op) {
+    return ::mlir::OpTrait::impl::verifyNoRegionArguments(op);
+  }
+};
+
 } // end namespace OpTrait
 
 //===----------------------------------------------------------------------===//

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

@@ -1810,13 +1810,6 @@ static LogicalResult verify(AffineIfOp op) {
                                            condition.getNumDims())))
     return failure();
 
-  // Verify that the entry of each child region does not have arguments.
-  for (auto &region : op.getOperation()->getRegions()) {
-    for (auto &b : region)
-      if (b.getNumArguments() != 0)
-        return op.emitOpError(
-            "requires that child entry blocks have no arguments");
-  }
   return success();
 }
 

mlir/lib/Dialect/SCF/SCF.cpp

@@ -410,16 +410,6 @@ void IfOp::build(OpBuilder &builder, OperationState &result, Value cond,
 }
 
 static LogicalResult verify(IfOp op) {
-  // Verify that the entry of each child region does not have arguments.
-  for (auto &region : op.getOperation()->getRegions()) {
-    if (region.empty())
-      continue;
-
-    for (auto &b : region)
-      if (b.getNumArguments() != 0)
-        return op.emitOpError(
-            "requires that child entry blocks have no arguments");
-  }
   if (op.getNumResults() != 0 && op.elseRegion().empty())
     return op.emitOpError("must have an else block if defining values");
 

mlir/lib/IR/Module.cpp

@@ -76,11 +76,6 @@ LogicalResult ModuleOp::verify() {
   if (!llvm::hasSingleElement(bodyRegion))
     return emitOpError("expected body region to have a single block");
 
-  // Check that the body has no block arguments.
-  auto *body = &bodyRegion.front();
-  if (body->getNumArguments() != 0)
-    return emitOpError("expected body to have no arguments");
-
   // Check that none of the attributes are non-dialect attributes, except for
   // the symbol related attributes.
   for (auto attr : getOperation()->getMutableAttrDict().getAttrs()) {

mlir/lib/IR/Operation.cpp

@@ -1017,6 +1017,22 @@ LogicalResult OpTrait::impl::verifyResultSizeAttr(Operation *op,
   return verifyValueSizeAttr(op, attrName, /*isOperand=*/false);
 }
 
+LogicalResult OpTrait::impl::verifyNoRegionArguments(Operation *op) {
+  for (Region &region : op->getRegions()) {
+    if (region.empty())
+      continue;
+
+    if (region.front().getNumArguments() != 0) {
+      if (op->getNumRegions() > 1)
+        return op->emitOpError("region #")
+               << region.getRegionNumber() << " should have no arguments";
+      else
+        return op->emitOpError("region should have no arguments");
+    }
+  }
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // BinaryOp implementation
 //===----------------------------------------------------------------------===//

mlir/test/Dialect/Affine/invalid.mlir

@@ -309,3 +309,32 @@ func @vector_store_vector_memref() {
   }
   return
 }
+
+// -----
+
+func @affine_if_with_then_region_args(%N: index) {
+  %c = constant 200 : index
+  %i = constant 20: index
+  // expected-error@+1 {{affine.if' op region #0 should have no arguments}}
+  affine.if affine_set<(i)[N] : (i - 2 >= 0, 4 - i >= 0)>(%i)[%c]  {
+    ^bb0(%arg:i32):
+      %w = affine.apply affine_map<(d0,d1)[s0] -> (d0+d1+s0)> (%i, %i) [%N]
+  }
+  return
+}
+
+// -----
+
+func @affine_if_with_else_region_args(%N: index) {
+  %c = constant 200 : index
+  %i = constant 20: index
+  // expected-error@+1 {{affine.if' op region #1 should have no arguments}}
+  affine.if affine_set<(i)[N] : (i - 2 >= 0, 4 - i >= 0)>(%i)[%c]  {
+      %w = affine.apply affine_map<(d0,d1)[s0] -> (d0+d1+s0)> (%i, %i) [%N]
+  } else {
+    ^bb0(%arg:i32):
+      %w = affine.apply affine_map<(d0,d1)[s0] -> (d0-d1+s0)> (%i, %i) [%N]
+  }
+  return
+}
+

mlir/test/Dialect/SCF/invalid.mlir

@@ -105,7 +105,7 @@ func @loop_if_not_one_block_per_region(%arg0: i1) {
 // -----
 
 func @loop_if_illegal_block_argument(%arg0: i1) {
-  // expected-error@+1 {{requires that child entry blocks have no arguments}}
+  // expected-error@+1 {{region #0 should have no arguments}}
   "scf.if"(%arg0) ({
     ^bb0(%0 : index):
       scf.yield

mlir/test/IR/invalid-module-op.mlir

@@ -16,7 +16,7 @@ func @module_op() {
 // -----
 
 func @module_op() {
-  // expected-error@+1 {{expected body to have no arguments}}
+  // expected-error@+1 {{region should have no arguments}}
   module {
   ^bb1(%arg: i32):
     "module_terminator"() : () -> ()