[MLIR] Allow overriding type constraints for control-flow interfaces.

[arames]

Summary

The introduce method areCompatibleControlFlowEdgeOperandTypes on
BranchOpInterface and RegionBranchOpInterface interfaces allows specifying
type constraints between operand and block argument types.
In particular, it allows using the constraint isMoreSpecializedOrSame()
instead of equality.

[arames]

Rebase on top of tree.

mlir/include/mlir/Interfaces/ControlFlowInterfaces.td

let methods = [
     >,
     InterfaceMethod<[{
         Returns the successor that would be chosen with the given constant
         operands. Returns nullptr if a single successor could not be chosen.
       }],
       "::mlir::Block *", "getSuccessorForOperands",
       (ins "::mlir::ArrayRef<::mlir::Attribute>":$operands), [{}],
       /*defaultImplementation=*/[{ return nullptr; }]
+    >,
+    StaticInterfaceMethod<
+      /*desc=*/"Specify type constraints over control-flow edges.",
+      /*retTy=*/"bool",
+      /*methodName=*/"areCompatibleControlFlowEdgeOperandTypes",
+      /*args=*/(ins "::mlir::Type":$operandTy, "::mlir::Type":$blockArgTy),
+      /*methodBody=*/[{}],
+      /*defaultImplementation=*/"return operandTy == blockArgTy;"
       >
   ];
 
   let verify = [{
     auto concreteOp = ::mlir::cast<ConcreteOp>($_op);
     for (unsigned i = 0, e = $_op->getNumSuccessors(); i != e; ++i) {
       ::mlir::Optional<OperandRange> operands = concreteOp.getSuccessorOperands(i);
       if (::mlir::failed(::mlir::detail::verifyBranchSuccessorOperands($_op, i, operands)))
         return ::mlir::failure();
InterfaceMethod<[{
       "void", "getNumRegionInvocations",
       (ins "::mlir::ArrayRef<::mlir::Attribute>":$operands,
            "::mlir::SmallVectorImpl<int64_t> &":$countPerRegion), [{}],
       /*defaultImplementation=*/[{
         unsigned numRegions = this->getOperation()->getNumRegions();
         assert(countPerRegion.empty());
         countPerRegion.resize(numRegions, kUnknownNumRegionInvocations);
       }]
+    >,
+    StaticInterfaceMethod<
+      /*desc=*/"Specify type constraints over control-flow edges.",
+      /*retTy=*/"bool",
+      /*methodName=*/"areCompatibleControlFlowEdgeOperandTypes",
+      /*args=*/(ins "::mlir::Type":$operandTy, "::mlir::Type":$blockArgTy),
+      /*methodBody=*/[{}],
+      /*defaultImplementation=*/"return operandTy == blockArgTy;"
       >
   ];
 
   let verify = [{
     static_assert(!ConcreteOp::template hasTrait<OpTrait::ZeroRegion>(),
                   "expected operation to have non-zero regions");
     return success();
   }];
 

mlir/lib/Interfaces/ControlFlowInterfaces.cpp

if (operandCount != destBB->getNumArguments())
     return op->emitError() << "branch has " << operandCount
                            << " operands for successor #" << succNo
                            << ", but target block has "
                            << destBB->getNumArguments();
 
   // Check the types.
   auto operandIt = operands->begin();
   for (unsigned i = 0; i != operandCount; ++i, ++operandIt) {
-    if ((*operandIt).getType() != destBB->getArgument(i).getType())
+    if (!cast<BranchOpInterface>(op).areCompatibleControlFlowEdgeOperandTypes(
+            (*operandIt).getType(), destBB->getArgument(i).getType()))
       return op->emitError() << "type mismatch for bb argument #" << i
                              << " of successor #" << succNo;
   }
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // RegionBranchOpInterface
if (sourceTypes->size() != succInputsTypes.size()) {
                                  << " operands, but target successor needs "
                                  << succInputsTypes.size();
     }
 
     for (auto typesIdx :
          llvm::enumerate(llvm::zip(*sourceTypes, succInputsTypes))) {
       Type sourceType = std::get<0>(typesIdx.value());
       Type inputType = std::get<1>(typesIdx.value());
-      if (sourceType != inputType) {
+      if (!regionInterface.areCompatibleControlFlowEdgeOperandTypes(
+              sourceType, inputType)) {
         InFlightDiagnostic diag = op->emitOpError(" along control flow edge ");
         return printEdgeName(diag)
                << ": source type #" << typesIdx.index() << " " << sourceType
                << " should match input type #" << typesIdx.index() << " "
                << inputType;
       }
     }
   }

mlir/test/IR/invalid.mlir

 ^bb1(%a: i64):  // expected-note {{previously defined here}}
   br ^bb2
 ^bb2(%a: i64):  // expected-error{{redefinition of SSA value '%a'}}
   return
 }
 
 // -----
 
-func @br_mismatch() {
-^bb0:
-  %0:2 = "foo"() : () -> (i1, i17)
-  // expected-error @+1 {{branch has 2 operands for successor #0, but target block has 1}}
-  br ^bb1(%0#1, %0#0 : i17, i1)
-
-^bb1(%x: i17):
-  return
-}
-
-// -----
-
-func @succ_arg_type_mismatch() {
-^bb0:
-  %0 = "getBool"() : () -> i1
-  // expected-error @+1 {{type mismatch for bb argument #0 of successor #0}}
-  br ^bb1(%0 : i1)
-
-^bb1(%x: i32):
-  return
-}
-
-
-// -----
-
 // Test no nested vector.
 func @vectors(vector<1 x vector<1xi32>>, vector<2x4xf32>)
 // expected-error@-1 {{vector elements must be int/index/float type}}
 
 // -----
 
 func @condbr_notbool() {
 ^bb0:

mlir/test/Interfaces/ControlFlowInterfaces/operand-types.mlir

+// RUN: mlir-opt -allow-unregistered-dialect --split-input-file --verify-diagnostics %s
+
+//===----------------------------------------------------------------------===//
+// Test BranchOpInterface
+//===----------------------------------------------------------------------===//
+
+func @succ_arg_number_mismatch() {
+^bb0:
+  %values:2 = "getValues"() : () -> (i1, i17)
+  "test.br"(%values#1, %values#0)[^bb1] : (i17, i1) -> ()
+  // expected-error@-1 {{branch has 2 operands for successor #0, but target block has 1}}
+
+^bb1(%arg1: i17):
+  return
+}
+
+// -----
+
+func @succ_arg_type_mismatch() {
+^bb0:
+  %value = "getValue"() : () -> i1
+  "test.br"(%value)[^bb1] : (i1) -> ()
+  // expected-error@-1 {{type mismatch for bb argument #0 of successor #0}}
+
+^bb1(%arg1: i32):
+  return
+}
+
+// -----
+
+func @succ_type_fixed_to_dynamic() {
+^bb0:
+  %0 = "getValue"() : () -> tensor<1xi32>
+  "test.br"(%0)[^bb1] : (tensor<1xi32>) -> ()
+  // expected-error@-1 {{type mismatch for bb argument #0 of successor #0}}
+^bb1(%arg1: tensor<?xi32>):
+  return
+
+^bb2:
+  "test.cfe.br"(%0)[^bb3] : (tensor<1xi32>) -> ()
+^bb3(%arg2: tensor<?xi32>):
+  return
+}
+
+// -----
+
+func @succ_type_dynamic_to_fixed() {
+^bb0:
+  %0 = "getValue"() : () -> tensor<?xi32>
+  "test.cfe.br"(%0)[^bb1] : (tensor<?xi32>) -> ()
+  // expected-error@-1 {{type mismatch for bb argument #0 of successor #0}}
+^bb1(%arg3: tensor<1xi32>):
+  return
+}
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// Test RegionBranchOpInterface
+//===----------------------------------------------------------------------===//
+
+func @succ_arg_type_match() {
+  %0 = "getValue"() : () -> (i32)
+
+  %1 = test.region_if %0 : i32 -> (i32) then {
+    ^bb0(%arg1 : i32):
+      test.region_if_yield %arg1 : i32
+  } else {
+    ^bb0(%arg1 : i32):
+      test.region_if_yield %arg1 : i32
+  } join {
+    ^bb0(%arg1 : i32):
+      test.region_if_yield %arg1 : i32
+  }
+
+  return
+}
+
+// -----
+
+func @succ_type_fixed_to_dynamic() {
+  %0 = "getValue"() : () -> (memref<1xi32>)
+
+  // expected-error@+1 {{'test.region_if' op  along control flow edge from parent operands to Region #1: source type #0 'memref<1xi32>' should match input type #0 'memref<?xi32>'}}
+  %tmp1 = test.region_if %0 : memref<1xi32> -> (memref<?xi32>) then {
+    ^bb0(%arg1 : memref<1xi32>):
+      %true_value = "getValue"(%arg1) : (memref<1xi32>) -> (memref<2xi32>)
+      test.region_if_yield %true_value : memref<2xi32>
+  } else {
+    ^bb0(%arg1 : memref<?xi32>):
+      %false_value = "getValue"(%arg1) : (memref<?xi32>) -> (memref<3xi32>)
+      test.region_if_yield %false_value : memref<3xi32>
+  } join {
+    ^bb0(%arg1 : memref<?xi32>):
+      test.region_if_yield %arg1 : memref<?xi32>
+  }
+
+  %tmp2 = test.cfe.region_if %0 : memref<1xi32> -> (memref<?xi32>) then {
+    ^bb0(%arg1 : memref<1xi32>):
+      %true_value = "getValue"(%arg1) : (memref<1xi32>) -> (memref<2xi32>)
+      test.cfe.region_if_yield %true_value : memref<2xi32>
+  } else {
+    ^bb0(%arg1 : memref<?xi32>):
+      %false_value = "getValue"(%arg1) : (memref<?xi32>) -> (memref<3xi32>)
+      test.cfe.region_if_yield %false_value : memref<3xi32>
+  } join {
+    ^bb0(%arg1 : memref<?xi32>):
+      test.cfe.region_if_yield %arg1 : memref<?xi32>
+  }
+
+  return
+}

mlir/test/lib/Dialect/Test/TestDialect.cpp

if (auto str = value[i].dyn_cast<StringAttr>())
       if (!str.getValue().empty())
         setNameFn(getResult(i), str.getValue());
 }
 
 //===----------------------------------------------------------------------===//
 // RegionIfOp
 //===----------------------------------------------------------------------===//
 
-static void print(OpAsmPrinter &p, RegionIfOp op) {
+template <typename RegionIfOpT>
+static void printRegionIfOp(OpAsmPrinter &p, RegionIfOpT op) {
   p << " ";
   p.printOperands(op.getOperands());
   p << ": " << op.getOperandTypes();
   p.printArrowTypeList(op.getResultTypes());
   p << " then";
   p.printRegion(op.getThenRegion(),
                 /*printEntryBlockArgs=*/true,
                 /*printBlockTerminators=*/true);
   p << " else";
   p.printRegion(op.getElseRegion(),
                 /*printEntryBlockArgs=*/true,
                 /*printBlockTerminators=*/true);
   p << " join";
   p.printRegion(op.getJoinRegion(),
                 /*printEntryBlockArgs=*/true,
                 /*printBlockTerminators=*/true);
 }
 
+static void print(OpAsmPrinter &p, RegionIfOp op) { printRegionIfOp(p, op); }
+
 static ParseResult parseRegionIfOp(OpAsmParser &parser,
                                    OperationState &result) {
   SmallVector<OpAsmParser::OperandType, 2> operandInfos;
   SmallVector<Type, 2> operandTypes;
 
   result.regions.reserve(3);
   Region *thenRegion = result.addRegion();
   Region *elseRegion = result.addRegion();

 static void print(SingleNoTerminatorCustomAsmOp op, OpAsmPrinter &printer) {
   printer.printRegion(
       op.getRegion(), /*printEntryBlockArgs=*/false,
       // This op has a single block without terminators. But explicitly mark
       // as not printing block terminators for testing.
       /*printBlockTerminators=*/false);
 }
 
+//===----------------------------------------------------------------------===//
+// Control-Flow Edge Operand Types Test Ops
+//===----------------------------------------------------------------------===//
+
+Optional<MutableOperandRange>
+CFEBranchOp::getMutableSuccessorOperands(unsigned index) {
+  assert(index == 0 && "invalid successor index");
+  return getTargetOperandsMutable();
+}
+
+bool CFEBranchOp::areCompatibleControlFlowEdgeOperandTypes(Type operandTy,
+                                                           Type blockArgTy) {
+  return isMoreSpecializedOrSame(operandTy, blockArgTy);
+}
+
+static void print(OpAsmPrinter &p, CFERegionIfOp op) { printRegionIfOp(p, op); }
+
+OperandRange CFERegionIfOp::getSuccessorEntryOperands(unsigned index) {
+  assert(index < 2 && "invalid region index");
+  return getOperands();
+}
+
+void CFERegionIfOp::getSuccessorRegions(
+    Optional<unsigned> index, ArrayRef<Attribute> operands,
+    SmallVectorImpl<RegionSuccessor> &regions) {
+  // We always branch to the join region.
+  if (index.hasValue()) {
+    if (index.getValue() < 2)
+      regions.push_back(RegionSuccessor(&getJoinRegion(), getJoinArgs()));
+    else
+      regions.push_back(RegionSuccessor(getResults()));
+    return;
+  }
+
+  // The then and else regions are the entry regions of this op.
+  regions.push_back(RegionSuccessor(&getThenRegion(), getThenArgs()));
+  regions.push_back(RegionSuccessor(&getElseRegion(), getElseArgs()));
+}
+
+bool CFERegionIfOp::areCompatibleControlFlowEdgeOperandTypes(Type operandTy,
+                                                             Type blockArgTy) {
+  return isMoreSpecializedOrSame(operandTy, blockArgTy);
+}
+
 #include "TestOpEnums.cpp.inc"
 #include "TestOpInterfaces.cpp.inc"
 #include "TestOpStructs.cpp.inc"
 #include "TestTypeInterfaces.cpp.inc"
 
 #define GET_OP_CLASSES
 #include "TestOps.cpp.inc"

mlir/test/lib/Dialect/Test/TestOps.td

let extraClassDeclaration = [{
     ::mlir::Block::BlockArgListType getElseArgs() {
       return getBody(1)->getArguments();
     }
     ::mlir::Block::BlockArgListType getJoinArgs() {
       return getBody(2)->getArguments();
     }
     ::mlir::OperandRange getSuccessorEntryOperands(unsigned index);
   }];
+
+  let verifier = [{ return RegionBranchOpInterface::verifyTypes(*this); }];
 }
 
 //===----------------------------------------------------------------------===//
 // Test TableGen generated build() methods
 //===----------------------------------------------------------------------===//
 
 def TableGenConstant : TEST_Op<"tblgen_constant"> {
   let results = (outs AnyType);

 def TestDefaultStrAttrHasValueOp : TEST_Op<"has_str_value"> {
   let arguments = (ins DefaultValuedStrAttr<StrAttr, "">:$value);
   let assemblyFormat = "attr-dict";
 }
 
 def : Pat<(TestDefaultStrAttrNoValueOp $value),
           (TestDefaultStrAttrHasValueOp ConstantStrAttr<StrAttr, "foo">)>;
 
+//===----------------------------------------------------------------------===//
+// Test Control-Flow Edge Operand Types
+//===----------------------------------------------------------------------===//
+
+def CFEBranchOp : TEST_Op<"cfe.br",
+    [DeclareOpInterfaceMethods<BranchOpInterface,
+                               ["areCompatibleControlFlowEdgeOperandTypes"]>,
+     Terminator]> {
+  let arguments = (ins Variadic<AnyType>:$targetOperands);
+  let successors = (successor AnySuccessor:$target);
+}
+
+def CFERegionIfYieldOp : TEST_Op<"cfe.region_if_yield",
+      [NoSideEffect, ReturnLike, Terminator]> {
+  let arguments = (ins Variadic<AnyType>:$results);
+  let assemblyFormat = [{
+    $results `:` type($results) attr-dict
+  }];
+}
+
+def CFERegionIfOp
+    : TEST_Op<"cfe.region_if",
+      [DeclareOpInterfaceMethods<RegionBranchOpInterface,
+                                 ["areCompatibleControlFlowEdgeOperandTypes"]>,
+       SingleBlockImplicitTerminator<"CFERegionIfYieldOp">,
+       RecursiveSideEffects]> {
+  let description =[{
+    Represents an abstract if-then-else-join pattern. In this context, the then
+    and else regions jump to the join region, which finally returns to its
+    parent op.
+    }];
+
+  let printer = [{ return ::print(p, *this); }];
+  let parser = [{ return ::parseRegionIfOp(parser, result); }];
+  let arguments = (ins Variadic<AnyType>);
+  let results = (outs Variadic<AnyType>:$results);
+  let regions = (region SizedRegion<1>:$thenRegion,
+                        AnyRegion:$elseRegion,
+                        AnyRegion:$joinRegion);
+  let extraClassDeclaration = [{
+    ::mlir::Block::BlockArgListType getThenArgs() {
+      return getBody(0)->getArguments();
+    }
+    ::mlir::Block::BlockArgListType getElseArgs() {
+      return getBody(1)->getArguments();
+    }
+    ::mlir::Block::BlockArgListType getJoinArgs() {
+      return getBody(2)->getArguments();
+    }
+    ::mlir::OperandRange getSuccessorEntryOperands(unsigned index);
+  }];
+
+  let verifier = [{ return RegionBranchOpInterface::verifyTypes(*this); }];
+}
+
 #endif // TEST_OPS