[MLIR] Add a NoRegionArguments trait

[jurahul]

Summary

• 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

[rriddle]

Did you consider using a per-region constraint in ODS instead?

mlir/lib/IR/Operation.cpp

1025

nit: Can you merge this condition into the previous else?

1026

nit: Add braces here as this body isn't "trivial".

1029

nit: Drop else after a return.

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

let extraClassDeclaration = [{
     bool matchingBoundOperandList();
   }];
 
   let hasCanonicalizer = 1;
   let hasFolder = 1;
 }
 
 def AffineIfOp : Affine_Op<"if",
-                           [ImplicitAffineTerminator, RecursiveSideEffects]> {
+                           [ImplicitAffineTerminator, RecursiveSideEffects,
+                            NoRegionArguments]> {
   let summary = "if-then-else operation";
   let description = [{
     Syntax:
 
     ```
     operation  ::= `affine.if` if-op-cond `{` op* `}` (`else` `{` op* `}`)?
     if-op-cond ::= integer-set-attr dim-and-symbol-use-list
     ```

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

let extraClassDeclaration = [{
     /// induction variable. LoopOp only has one region, so 0 is the only valid
     /// value for `index`.
     OperandRange getSuccessorEntryOperands(unsigned index);
   }];
 }
 
 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
     conditionally executing two regions of code. The operand to an if operation
     is a boolean value. For example:
 
     ```mlir
     scf.if %b  {

mlir/include/mlir/IR/Module.h

 /// allowed to implicitly capture global values, and all external references
 /// must use symbolic references via attributes(e.g. via a string name).
 class ModuleOp
     : public Op<
           ModuleOp, OpTrait::ZeroOperands, OpTrait::ZeroResult,
           OpTrait::IsIsolatedFromAbove, OpTrait::AffineScope,
           OpTrait::SymbolTable,
           OpTrait::SingleBlockImplicitTerminator<ModuleTerminatorOp>::Impl,
-          SymbolOpInterface::Trait> {
+          SymbolOpInterface::Trait, OpTrait::NoRegionArguments> {
 public:
   using Op::Op;
   using Op::print;
 
   static StringRef getOperationName() { return "module"; }
 
   static void build(OpBuilder &builder, OperationState &result,
                     Optional<StringRef> name = llvm::None);

mlir/include/mlir/IR/OpBase.td

 // vector that has the same number of elements as the number of ODS declared
 // operands. That means even if some operands are non-variadic, the attribute
 // still need to have an element for its size, which is always 1.
 def AttrSizedOperandSegments : NativeOpTrait<"AttrSizedOperandSegments">;
 // Similar to AttrSizedOperandSegments, but used for results. The attribute
 // should be named as `result_segment_sizes`.
 def AttrSizedResultSegments  : NativeOpTrait<"AttrSizedResultSegments">;
 
+// Op attached regions have no arguments
+def NoRegionArguments : NativeOpTrait<"NoRegionArguments">;
+
 //===----------------------------------------------------------------------===//
 // OpInterface definitions
 //===----------------------------------------------------------------------===//
 
 // Marker used to identify the argument list for an op or interface method.
 def ins;
 
 // OpInterfaceTrait corresponds to a specific 'OpInterface' class defined in

mlir/include/mlir/IR/OpDefinition.h

 LogicalResult verifyResultsAreSignlessIntegerLike(Operation *op);
 LogicalResult verifyIsTerminator(Operation *op);
 LogicalResult verifyZeroSuccessor(Operation *op);
 LogicalResult verifyOneSuccessor(Operation *op);
 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
 /// with this directly, so its members are all protected.
 template <typename ConcreteType, template <typename> class TraitType>
 class TraitBase {
 protected:
   /// Return the ultimate Operation being worked on.
public:
   static StringRef getResultSegmentSizeAttr() { return "result_segment_sizes"; }
 
   static LogicalResult verifyTrait(Operation *op) {
     return ::mlir::OpTrait::impl::verifyResultSizeAttr(
         op, getResultSegmentSizeAttr());
   }
 };
 
+/// 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
 
 //===----------------------------------------------------------------------===//
 // Operation Definition classes
 //===----------------------------------------------------------------------===//
 
 /// This provides public APIs that all operations should have.  The template
 /// argument 'ConcreteType' should be the concrete type by CRTP and the others

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

return op.emitOpError(
         "operand count and condition integer set dimension and "
         "symbol count must match");
 
   // Verify that the operands are valid dimension/symbols.
   if (failed(verifyDimAndSymbolIdentifiers(op, op.getOperands(),
                                            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();
 }
 
 static ParseResult parseAffineIfOp(OpAsmParser &parser,
                                    OperationState &result) {
   // Parse the condition attribute set.
   IntegerSetAttr conditionAttr;
   unsigned numDims;

mlir/lib/Dialect/SCF/SCF.cpp

 
 void IfOp::build(OpBuilder &builder, OperationState &result, Value cond,
                  function_ref<void(OpBuilder &, Location)> thenBuilder,
                  function_ref<void(OpBuilder &, Location)> elseBuilder) {
   build(builder, result, TypeRange(), cond, thenBuilder, elseBuilder);
 }
 
 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");
 
   return success();
 }
 
 static ParseResult parseIfOp(OpAsmParser &parser, OperationState &result) {
   // Create the regions for 'then'.

mlir/lib/IR/Module.cpp

 
 LogicalResult ModuleOp::verify() {
   auto &bodyRegion = getOperation()->getRegion(0);
 
   // The body must contain a single basic block.
   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()) {
     if (!attr.first.strref().contains('.') &&
         !llvm::is_contained(
             ArrayRef<StringRef>{mlir::SymbolTable::getSymbolAttrName(),
                                 mlir::SymbolTable::getVisibilityAttrName()},
             attr.first.strref()))

mlir/lib/IR/Operation.cpp

LogicalResult OpTrait::impl::verifyOperandSizeAttr(Operation *op,
   return verifyValueSizeAttr(op, attrName, /*isOperand=*/true);
 }
 
 LogicalResult OpTrait::impl::verifyResultSizeAttr(Operation *op,
                                                   StringRef attrName) {
   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
 //===----------------------------------------------------------------------===//
 
 // These functions are out-of-line implementations of the methods in BinaryOp,
 // which avoids them being template instantiated/duplicated.
 
 void impl::buildBinaryOp(OpBuilder &builder, OperationState &result, Value lhs,

mlir/test/Dialect/Affine/invalid.mlir

func @vector_store_vector_memref() {
   %0 = alloc() : memref<100xvector<8xf32>>
   %1 = constant dense<7.0> : vector<8xf32>
   affine.for %i0 = 0 to 4 {
     // expected-error@+1 {{requires memref and vector types of the same elemental type}}
     affine.vector_store %1, %0[%i0] : memref<100xvector<8xf32>>, vector<8xf32>
   }
   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

^bb1:
       scf.yield
   }, {}): (i1) -> ()
   return
 }
 
 // -----
 
 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
   }, {}): (i1) -> ()
   return
 }
 
 // -----

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

^bb2:
     "module_terminator"() : () -> ()
   }
   return
 }
 
 // -----
 
 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"() : () -> ()
   }
   return
 }
 
 // -----