[mlir][interfaces] Add insideMutuallyExclusiveRegions helper

[springerm]

Summary

Add a helper function to ControlFlowInterfaces for checking if two ops are in mutually exclusive regions according to RegionBranchOpInterface.

Utilize this new helper in Linalg ComprehensiveBufferize. This makes the analysis independent of the SCF dialect and generalizes it to other ops that implement RegionBranchOpInterface.

Depends On D114219

[nicolasvasilache]

Could you add a test with a special op that has similar behavior to a loopy-switch (e.g. round-robin) implemented with the RegionBranchOpInterface and make sure you get what you expect?

[springerm]

move insideMutuallyExclusiveRegions to ControlFlowInterfaces

[ftynse]

mlir/lib/Interfaces/ControlFlowInterfaces.cpp

238 ↗

(On Diff #389637)

Nit: could you swap the condition and use continue here?

264 ↗

(On Diff #389637)

I remember there is a lifetime-related issue with llvm::enumerate and references so please use auto it here instead.

[springerm]

address comments

mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp

do {
     if (a->isProperAncestor(b))
       return false;
     if (domInfo.properlyDominates(a, b))
       return true;
   } while ((a = a->getParentOp()));
   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
+///    RegionBranchOpInterface::getSuccessorRegions.
+static bool insideMutuallyExclusiveRegions(Operation *a, Operation *b) {
+  assert(a && "expected non-empty operation");
+  assert(b && "expected non-empty operation");
+
+  auto branchOp = a->getParentOfType<RegionBranchOpInterface>();
+  while (branchOp) {
+    // Check if b is inside branchOp. (We already know that a is.)
+    if (branchOp->isProperAncestor(b)) {
+      // b is contained in branchOp. Retrieve the regions in which `a` and `b`
+      // are contained.
+      Region *regionA = nullptr, *regionB = nullptr;
+      for (Region &r : branchOp->getRegions()) {
+        if (r.findAncestorOpInRegion(*a)) {
+          assert(!regionA && "already found a region for a");
+          regionA = &r;
+        }
+        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) &&
+             !isRegionReachable(regionB, regionA);
+    }
+
+    // Check next enclosing RegionBranchOpInterface.
+    branchOp = branchOp->getParentOfType<RegionBranchOpInterface>();
+  }
+
+  // Could not find a common RegionBranchOpInterface among a's and b's
+  // ancestors.
+  return false;
+}
+
 /// Given sets of uses and writes, return true if there is a RaW conflict under
 /// the assumption that all given reads/writes alias the same buffer and that
 /// all given writes bufferize inplace.
 ///
 /// A conflict is: According to SSA use-def chains, a read R is supposed to read
 /// the result of a write W1. But because of bufferization decisions, R actually
 /// reads another write W2.
 static bool
for (OpOperand *uConflictingWrite : usesWrite) {
 
       if (conflictingWritingOp != readingOp)
         if (auto bufferizableOp =
                 dyn_cast<BufferizableOpInterface>(conflictingWritingOp))
           if (bufferizableOp.isNotConflicting(uRead, uConflictingWrite,
                                               aliasInfo))
             continue;
 
-      // Special rules for branches.
-      // TODO: Use an interface.
-      if (scf::insideMutuallyExclusiveBranches(readingOp, conflictingWritingOp))
+      // Ops are not conflicting if they are in mutually exclusive regions.
+      if (insideMutuallyExclusiveRegions(readingOp, conflictingWritingOp))
         continue;
 
       LDBG("WRITE = #" << printValueInfo(lastWrite) << "\n");
 
       // No conflict if the conflicting write happens before the last
       // write.
       if (Operation *writingOp = lastWrite.getDefiningOp()) {
         if (happensBefore(conflictingWritingOp, writingOp, domInfo))