[mlir] allow region branch spec from parent op to itself

[ftynse]

Summary

RegionBranchOpInterface did not allow the operation with regions to
specify itself as successors. Therefore, this implied that the control
is always transferred to a region before being transferred back to the
parent op. Since the region can only transfer the control back to the
parent op from a terminator, this transitively implied that the first
block of any region with a RegionBranchOpInterface is always executed
until the terminator can transfer the control flow back. This is
trivially false for any conditional-like operation that may or may not
execute the region, as well as for loop-like operations that may not
execute the body.

Remove the restriction from the interface description and update the
only transform that relied on it.

See
https://discourse.llvm.org/t/rfc-region-control-flow-interfaces-should-encode-region-not-executed-correctly/72103.

Depends On: https://reviews.llvm.org/D155757

mlir/include/mlir/Interfaces/ControlFlowInterfaces.td

let methods = [
     InterfaceMethod<[{
         Returns the viable successors of a region at `index`, or the possible
         successors when branching from the parent op if `index` is None. These
         are the regions that may be selected during the flow of control. If
         `index` is None, `operands` is a set of optional attributes that
         either correspond to a constant value for each operand of this
         operation, or null if that operand is not a constant. If `index` is
         valid, `operands` corresponds to the entry values of the region at
-        `index`. Only a region, i.e. a valid `index`, may use the parent
-        operation as a successor. This method allows for describing which
+        `index`. The parent operation, i.e. a null `index`, may specify itself
+        as successor, which indicates that the control flow may not enter any
+        region at all. This method allows for describing which
         regions may be executed when entering an operation, and which regions
         are executed after having executed another region of the parent op. The
         successor region must be non-empty.
       }],
       "void", "getSuccessorRegions",
       (ins "::std::optional<unsigned>":$index,
            "::llvm::ArrayRef<::mlir::Attribute>":$operands,
            "::llvm::SmallVectorImpl<::mlir::RegionSuccessor> &":$regions)

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

void AffineIfOp::getSuccessorRegions(
     std::optional<unsigned> index, ArrayRef<Attribute> operands,
     SmallVectorImpl<RegionSuccessor> &regions) {
   // If the predecessor is an AffineIfOp, then branching into both `then` and
   // `else` region is valid.
   if (!index.has_value()) {
     regions.reserve(2);
     regions.push_back(
         RegionSuccessor(&getThenRegion(), getThenRegion().getArguments()));
-    // Don't consider the else region if it is empty.
-    if (!getElseRegion().empty())
+    // If the "else" region is empty, branch bach into parent.
+    if (getElseRegion().empty()) {
+      regions.push_back(getResults());
+    } else {
       regions.push_back(
           RegionSuccessor(&getElseRegion(), getElseRegion().getArguments()));
+    }
     return;
   }
 
   // If the predecessor is the `else`/`then` region, then branching into parent
   // op is valid.
   regions.push_back(RegionSuccessor(getResults()));
 }
 

mlir/lib/Dialect/Bufferization/Transforms/BufferViewFlowAnalysis.cpp

if (auto regionInterface = dyn_cast<RegionBranchOpInterface>(op)) {
       // Query the RegionBranchOpInterface to find potential successor regions.
       // Extract all entry regions and wire all initial entry successor inputs.
       SmallVector<RegionSuccessor, 2> entrySuccessors;
       regionInterface.getSuccessorRegions(/*index=*/std::nullopt,
                                           entrySuccessors);
       for (RegionSuccessor &entrySuccessor : entrySuccessors) {
         // Wire the entry region's successor arguments with the initial
         // successor inputs.
-        assert(entrySuccessor.getSuccessor() &&
-               "Invalid entry region without an attached successor region");
         registerDependencies(
             regionInterface.getSuccessorEntryOperands(
-                entrySuccessor.getSuccessor()->getRegionNumber()),
+                entrySuccessor.isParent()
+                    ? std::optional<unsigned>()
+                    : entrySuccessor.getSuccessor()->getRegionNumber()),
             entrySuccessor.getSuccessorInputs());
       }
 
       // Wire flow between regions and from region exits.
       for (Region &region : regionInterface->getRegions()) {
         // Iterate over all successor region entries that are reachable from the
         // current region.
         SmallVector<RegionSuccessor, 2> successorRegions;

mlir/lib/Dialect/SCF/IR/SCF.cpp

if (!ivArg)
     return ForOp();
   assert(ivArg.getOwner() && "unlinked block argument");
   auto *containingOp = ivArg.getOwner()->getParentOp();
   return dyn_cast_or_null<ForOp>(containingOp);
 }
 
 /// Return operands used when entering the region at 'index'. These operands
 /// correspond to the loop iterator operands, i.e., those excluding the
-/// induction variable. LoopOp only has one region, so 0 is the only valid value
-/// for `index`.
+/// induction variable.
 OperandRange ForOp::getSuccessorEntryOperands(std::optional<unsigned> index) {
-  assert(index && *index == 0 && "invalid region index");
-
-  // The initial operands map to the loop arguments after the induction
-  // variable.
   return getInitArgs();
 }
 
 /// Given the region at `index`, or the parent operation if `index` is None,
 /// return the successor regions. These are the regions that may be selected
 /// during the flow of control. `operands` is a set of optional attributes that
 /// correspond to a constant value for each operand, or null if that operand is
 /// not a constant.
 void ForOp::getSuccessorRegions(std::optional<unsigned> index,
                                 ArrayRef<Attribute> operands,
                                 SmallVectorImpl<RegionSuccessor> &regions) {
-  // If the predecessor is the ForOp, branch into the body using the iterator
-  // arguments.
-  if (!index) {
-    regions.push_back(RegionSuccessor(&getLoopBody(), getRegionIterArgs()));
-    return;
-  }
-
-  // Otherwise, the loop may branch back to itself or the parent operation.
-  assert(*index == 0 && "expected loop region");
+  // Both the operation itself and the region may be branching into the body or
+  // back into the operation itself. It is possible for loop not to enter the
+  // body.
   regions.push_back(RegionSuccessor(&getLoopBody(), getRegionIterArgs()));
   regions.push_back(RegionSuccessor(getResults()));
 }
 
 Region &ForallOp::getLoopBody() { return getRegion(); }
 
 /// Promotes the loop body of a forallOp to its containing block if it can be
 /// determined that the loop has a single iteration.

 /// Given the region at `index`, or the parent operation if `index` is None,
 /// return the successor regions. These are the regions that may be selected
 /// during the flow of control. `operands` is a set of optional attributes that
 /// correspond to a constant value for each operand, or null if that operand is
 /// not a constant.
 void ForallOp::getSuccessorRegions(std::optional<unsigned> index,
                                    ArrayRef<Attribute> operands,
                                    SmallVectorImpl<RegionSuccessor> &regions) {
-  // If the predecessor is ForallOp, branch into the body with empty arguments.
-  if (!index) {
-    regions.push_back(RegionSuccessor(&getRegion()));
-    return;
-  }
-
-  // Otherwise, the loop should branch back to the parent operation.
-  assert(*index == 0 && "expected loop region");
+  // Both the operation itself and the region may be branching into the body or
+  // back into the operation itself. It is possible for loop not to enter the
+  // body.
   regions.push_back(RegionSuccessor());
 }
 
 //===----------------------------------------------------------------------===//
 // InParallelOp
 //===----------------------------------------------------------------------===//
 
 // Build a InParallelOp with mixed static and dynamic entries.
void IfOp::getSuccessorRegions(std::optional<unsigned> index,
 
   // Otherwise, the successor is dependent on the condition.
   bool condition;
   if (auto condAttr = llvm::dyn_cast_or_null<IntegerAttr>(operands.front())) {
     condition = condAttr.getValue().isOne();
   } else {
     // If the condition isn't constant, both regions may be executed.
     regions.push_back(RegionSuccessor(&getThenRegion()));
-    // If the else region does not exist, it is not a viable successor.
+    // If the else region does not exist, it is not a viable successor, so the
+    // control will go back to this operation instead.
     if (elseRegion)
       regions.push_back(RegionSuccessor(elseRegion));
+    else
+      regions.push_back(RegionSuccessor());
     return;
   }
 
   // Add the successor regions using the condition.
   regions.push_back(RegionSuccessor(condition ? &getThenRegion() : elseRegion));
 }
 
 LogicalResult IfOp::fold(FoldAdaptor adaptor,

 /// Given the region at `index`, or the parent operation if `index` is None,
 /// return the successor regions. These are the regions that may be selected
 /// during the flow of control. `operands` is a set of optional attributes that
 /// correspond to a constant value for each operand, or null if that operand is
 /// not a constant.
 void ParallelOp::getSuccessorRegions(
     std::optional<unsigned> index, ArrayRef<Attribute> operands,
     SmallVectorImpl<RegionSuccessor> &regions) {
-  // If the predecessor is ParallelOp, branch into the body with empty
-  // arguments.
-  if (!index) {
-    regions.push_back(RegionSuccessor(&getRegion()));
-    return;
-  }
-
-  assert(*index == 0 && "expected loop region");
-  // Otherwise, the loop should branch back to the parent operation.
+  // Both the operation itself and the region may be branching into the body or
+  // back into the operation itself. It is possible for loop not to enter the
+  // body.
   regions.push_back(RegionSuccessor());
 }
 
 //===----------------------------------------------------------------------===//
 // ReduceOp
 //===----------------------------------------------------------------------===//
 
 void ReduceOp::build(

mlir/test/Analysis/DataFlow/test-next-access.mlir

 ^bb2(%phi: f32):
   return %phi : f32
 }
 
 // CHECK-LABEL: @loop
 func.func @loop(%arg0: memref<?xf32>, %arg1: f32, %arg2: index, %arg3: index, %arg4: index) -> f32 {
   %c0 = arith.constant 0.0 : f32
   // CHECK:      name = "pre"
-  // CHECK-SAME: next_access = {{\[}}["loop"], "unknown"]
-  // The above is not entirely correct when the loop has 0 iterations, but 
-  // the region control flow specificaiton is currently incapable of
-  // specifying that.
+  // CHECK-SAME: next_access = {{\[}}["outside", "loop"], "unknown"]
   memref.load %arg0[%arg4] {name = "pre"} : memref<?xf32>
   %l = scf.for %i = %arg2 to %arg3 step %arg4 iter_args(%ia = %c0) -> (f32) {
     // CHECK:      name = "loop"
     // CHECK-SAME: next_access = {{\[}}["outside", "loop"], "unknown"]
     %0 = memref.load %arg0[%i] {name = "loop"} : memref<?xf32>
     %1 = arith.addf %ia, %0 : f32
     scf.yield %1 : f32
   }
   %v = memref.load %arg0[%arg3] {name = "outside"} : memref<?xf32>
   %2 = arith.addf %v, %l : f32
   return %2 : f32
 }
 
 // CHECK-LABEL: @conditional
 func.func @conditional(%cond: i1, %arg0: memref<f32>) {
   // CHECK:      name = "pre"
-  // CHECK-SAME: next_access = {{\[}}["then"]]
-  // The above is not entirely correct when the condition is false, but 
-  // the region control flow specificaiton is currently incapable of
-  // specifying that.
+  // CHECK-SAME: next_access = {{\[}}["post", "then"]]
   memref.load %arg0[] {name = "pre"}: memref<f32>
   scf.if %cond {
     // CHECK:      name = "then"
     // CHECK-SAME: next_access = {{\[}}["post"]]
     memref.load %arg0[] {name = "then"} : memref<f32>
   }
   memref.load %arg0[] {name = "post"} : memref<f32>
   return
func.func @two_sided_conditional(%cond: i1, %arg0: memref<f32>) {
   memref.load %arg0[] {name = "post"} : memref<f32>
   return
 }
 
 // CHECK-LABEL: @dead_conditional
 func.func @dead_conditional(%arg0: memref<f32>) {
   %false = arith.constant 0 : i1
   // CHECK:      name = "pre"
-  // CHECK-SAME: next_access = ["unknown"]
-  // The above is not entirely correct when the condition is false, but 
-  // the region control flow specificaiton is currently incapable of
-  // specifying that.
+  // CHECK-SAME: next_access = {{\[}}["post"]]
   memref.load %arg0[] {name = "pre"}: memref<f32>
   scf.if %false {
     // CHECK:      name = "then"
     // CHECK-SAME: next_access = "not computed"
     memref.load %arg0[] {name = "then"} : memref<f32>
   }
   memref.load %arg0[] {name = "post"} : memref<f32>
   return
func.func @infinite_recursive_call(%arg0: memref<f32>) {
   return
 }
 
 // -----
 
 // CHECK-LABEL: @recursive_call
 func.func @recursive_call(%arg0: memref<f32>, %cond: i1) {
   // CHECK:      name = "pre"
-  // CHECK-SAME: next_access = {{\[}}["pre"]]
-  // The above is not entirely correct when the condition is false, but 
-  // the region control flow specificaiton is currently incapable of
-  // specifying that.
+  // CHECK-SAME: next_access = {{\[}}["post", "pre"]]
   memref.load %arg0[] {name = "pre"} : memref<f32>
   scf.if %cond {
     func.call @recursive_call(%arg0, %cond) : (memref<f32>, i1) -> ()
   }
   memref.load %arg0[] {name = "post"} : memref<f32>
   return
 }