[mlir][interfaces] Fix infinite loop in insideMutuallyExclusiveRegions

[springerm]

Summary

This function was missing a termination condition.

Depends On D123789

[springerm]

update

[springerm]

no change

[ftynse]

mlir/lib/Interfaces/ControlFlowInterfaces.cpp

242

Nit: you may want to assert that regions belong to the same op.

[springerm]

This has landed as a3005a406e195bf8e9798cdf4173fdadabac49bb.

mlir/lib/Interfaces/ControlFlowInterfaces.cpp

for (unsigned regionNo : llvm::seq(0U, op->getNumRegions())) {
 
     if (failed(verifyTypesAlongAllEdges(op, regionNo, inputTypesFromRegion)))
       return failure();
   }
 
   return success();
 }
 
+/// Return `true` if region `r` is reachable from region `begin` according to
+/// the RegionBranchOpInterface (by taking a branch).
+static bool isRegionReachable(Region *begin, Region *r) {
+  auto op = cast<RegionBranchOpInterface>(begin->getParentOp());
+  SmallVector<bool> visited(op->getNumRegions(), false);
+  visited[begin->getRegionNumber()] = true;
+
+  // Retrieve all successors of the region and enqueue them in the worklist.
+  SmallVector<unsigned> worklist;
+  auto enqueueAllSuccessors = [&](unsigned index) {
+    SmallVector<RegionSuccessor> successors;
+    op.getSuccessorRegions(index, successors);
+    for (RegionSuccessor successor : successors)
+      if (!successor.isParent())
+        worklist.push_back(successor.getSuccessor()->getRegionNumber());
+  };
+  enqueueAllSuccessors(begin->getRegionNumber());
+
+  // Process all regions in the worklist via DFS.
+  while (!worklist.empty()) {
+    unsigned nextRegion = worklist.pop_back_val();
+    if (nextRegion == r->getRegionNumber())
+      return true;
+    if (visited[nextRegion])
+      continue;
+    visited[nextRegion] = true;
+    enqueueAllSuccessors(nextRegion);
+  }
+
+  return false;
+}
+
 /// Return `true` if `a` and `b` are in mutually exclusive regions.
 ///
 /// 1. Find the first common of `a` and `b` (ancestor) that implements
 ///    RegionBranchOpInterface.
 /// 2. Determine the regions `regionA` and `regionB` in which `a` and `b` are
 ///    contained.
 /// 3. Check if `regionA` and `regionB` are mutually exclusive. They are
 ///    mutually exclusive if they are not reachable from each other as per
for (Region &r : branchOp->getRegions()) {
       }
       if (r.findAncestorOpInRegion(*b)) {
         assert(!regionB && "already found a region for b");
         regionB = &r;
       }
     }
     assert(regionA && regionB && "could not find region of op");
 
-    // Helper function that checks if region `r` is reachable from region
-    // `begin`.
-    std::function<bool(Region *, Region *)> isRegionReachable =
-        [&](Region *begin, Region *r) {
-          if (begin == r)
-            return true;
-          if (begin == nullptr)
-            return false;
-          // Compute index of region.
-          int64_t beginIndex = -1;
-          for (const auto &it : llvm::enumerate(branchOp->getRegions()))
-            if (&it.value() == begin)
-              beginIndex = it.index();
-          assert(beginIndex != -1 && "could not find region in op");
-          // Retrieve all successors of the region.
-          SmallVector<RegionSuccessor> successors;
-          branchOp.getSuccessorRegions(beginIndex, successors);
-          // Call function recursively on all successors.
-          for (RegionSuccessor successor : successors)
-            if (isRegionReachable(successor.getSuccessor(), r))
-              return true;
-          return false;
-        };
-
-    // `a` and `b` are in mutually exclusive regions if neither region is
-    // reachable from the other region.
-    return !isRegionReachable(regionA, regionB) &&
+    // `a` and `b` are in mutually exclusive regions if both regions are
+    // distinct and neither region is reachable from the other region.
+    return regionA != regionB && !isRegionReachable(regionA, regionB) &&
            !isRegionReachable(regionB, regionA);
   }
 
   // Could not find a common RegionBranchOpInterface among a's and b's
   // ancestors.
   return false;
 }
 
 bool RegionBranchOpInterface::isRepetitiveRegion(unsigned index) {
-  SmallVector<bool> visited(getOperation()->getNumRegions(), false);
-  visited[index] = true;
-
-  // Retrieve all successors of the region and enqueue them in the worklist.
-  SmallVector<unsigned> worklist;
-  auto enqueueAllSuccessors = [&](unsigned index) {
-    SmallVector<RegionSuccessor> successors;
-    this->getSuccessorRegions(index, successors);
-    for (RegionSuccessor successor : successors)
-      if (!successor.isParent())
-        worklist.push_back(successor.getSuccessor()->getRegionNumber());
-  };
-  enqueueAllSuccessors(index);
-
-  // Process all regions in the worklist via DFS.
-  while (!worklist.empty()) {
-    unsigned nextRegion = worklist.pop_back_val();
-    if (nextRegion == index)
-      return true;
-    if (visited[nextRegion])
-      continue;
-    visited[nextRegion] = true;
-    enqueueAllSuccessors(nextRegion);
-  }
-
-  return false;
+  Region *region = &getOperation()->getRegion(index);
+  return isRegionReachable(region, region);
 }
 
 Region *mlir::getEnclosingRepetitiveRegion(Operation *op) {
   while (Region *region = op->getParentRegion()) {
     op = region->getParentOp();
     if (auto branchOp = dyn_cast<RegionBranchOpInterface>(op))
       if (branchOp.isRepetitiveRegion(region->getRegionNumber()))
         return region;

mlir/unittests/Interfaces/ControlFlowInterfacesTest.cpp

if (index) {
         // This reguin also branches back to the parent.
         regions.push_back(RegionSuccessor());
       regions.push_back(
           RegionSuccessor(&getOperation()->getRegion(*index % kNumRegions)));
     }
   }
 };
 
+/// Each region branches back it itself or the parent.
+struct DoubleLoopRegionsOp
+    : public Op<DoubleLoopRegionsOp, RegionBranchOpInterface::Trait> {
+  using Op::Op;
+
+  static ArrayRef<StringRef> getAttributeNames() { return {}; }
+
+  static StringRef getOperationName() {
+    return "cftest.double_loop_regions_op";
+  }
+
+  void getSuccessorRegions(Optional<unsigned> index,
+                           ArrayRef<Attribute> operands,
+                           SmallVectorImpl<RegionSuccessor> &regions) {
+    if (index.hasValue()) {
+      regions.push_back(RegionSuccessor());
+      regions.push_back(RegionSuccessor(&getOperation()->getRegion(*index)));
+    }
+  }
+};
+
 /// Regions are executed sequentially.
 struct SequentialRegionsOp
     : public Op<SequentialRegionsOp, RegionBranchOpInterface::Trait> {
   using Op::Op;
   static ArrayRef<StringRef> getAttributeNames() { return {}; }
 
   static StringRef getOperationName() { return "cftest.sequential_regions_op"; }
 
   // Region 0 has Region 1 as a successor.
   void getSuccessorRegions(Optional<unsigned> index,
                            ArrayRef<Attribute> operands,
                            SmallVectorImpl<RegionSuccessor> &regions) {
     if (index == 0u) {
       Operation *thisOp = this->getOperation();
       regions.push_back(RegionSuccessor(&thisOp->getRegion(1)));
     }
   }
 };
 
 /// A dialect putting all the above together.
 struct CFTestDialect : Dialect {
   explicit CFTestDialect(MLIRContext *ctx)
       : Dialect(getDialectNamespace(), ctx, TypeID::get<CFTestDialect>()) {
     addOperations<DummyOp, MutuallyExclusiveRegionsOp, LoopRegionsOp,
-                  SequentialRegionsOp>();
+                  DoubleLoopRegionsOp, SequentialRegionsOp>();
   }
   static StringRef getDialectNamespace() { return "cftest"; }
 };
 
 TEST(RegionBranchOpInterface, MutuallyExclusiveOps) {
   const char *ir = R"MLIR(
 "cftest.mutually_exclusive_regions_op"() (
       {"cftest.dummy_op"() : () -> ()},  // op1
"cftest.mutually_exclusive_regions_op"() (
   Operation *testOp = &module->getBody()->getOperations().front();
   Operation *op1 = &testOp->getRegion(0).front().front();
   Operation *op2 = &testOp->getRegion(1).front().front();
 
   EXPECT_TRUE(insideMutuallyExclusiveRegions(op1, op2));
   EXPECT_TRUE(insideMutuallyExclusiveRegions(op2, op1));
 }
 
+TEST(RegionBranchOpInterface, MutuallyExclusiveOps2) {
+  const char *ir = R"MLIR(
+"cftest.double_loop_regions_op"() (
+      {"cftest.dummy_op"() : () -> ()},  // op1
+      {"cftest.dummy_op"() : () -> ()}   // op2
+  ) : () -> ()
+  )MLIR";
+
+  DialectRegistry registry;
+  registry.insert<CFTestDialect>();
+  MLIRContext ctx(registry);
+
+  OwningOpRef<ModuleOp> module = parseSourceString<ModuleOp>(ir, &ctx);
+  Operation *testOp = &module->getBody()->getOperations().front();
+  Operation *op1 = &testOp->getRegion(0).front().front();
+  Operation *op2 = &testOp->getRegion(1).front().front();
+
+  EXPECT_TRUE(insideMutuallyExclusiveRegions(op1, op2));
+  EXPECT_TRUE(insideMutuallyExclusiveRegions(op2, op1));
+}
+
 TEST(RegionBranchOpInterface, NotMutuallyExclusiveOps) {
   const char *ir = R"MLIR(
 "cftest.sequential_regions_op"() (
       {"cftest.dummy_op"() : () -> ()},  // op1
       {"cftest.dummy_op"() : () -> ()}   // op2
   ) : () -> ()
   )MLIR";