[mlir][IR] DominanceInfo: Add function to query dominator of a range of block

[matthias-springer]

Also improve the implementation of findCommonDominator (skip duplicate blocks) and extract it from BufferPlacementTransformationBase (so that BufferPlacementTransformationBase can be retired eventually).

[llvmbot]

@llvm/pr-subscribers-mlir
@llvm/pr-subscribers-mlir-core

@llvm/pr-subscribers-mlir-bufferization

Author: Matthias Springer (matthias-springer)

Changes

Also improve the implementation of findCommonDominator (skip duplicate blocks) and extract it from BufferPlacementTransformationBase (so that BufferPlacementTransformationBase can be retired eventually).

---

Full diff: https://github.com/llvm/llvm-project/pull/77098.diff

2 Files Affected:

• (modified) mlir/include/mlir/Dialect/Bufferization/Transforms/BufferUtils.h (+24-22)
• (modified) mlir/include/mlir/IR/Dominance.h (+15)

diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/BufferUtils.h b/mlir/include/mlir/Dialect/Bufferization/Transforms/BufferUtils.h
index 85e9c47ad5302c..26df4c655f2cda 100644
--- a/mlir/include/mlir/Dialect/Bufferization/Transforms/BufferUtils.h
+++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/BufferUtils.h
@@ -73,33 +73,35 @@ class BufferPlacementAllocs {
   AllocEntryList allocs;
 };
 
+/// Finds a common dominator for the given value while taking the positions
+/// of the values in the value set into account. It supports dominator and
+/// post-dominator analyses via template arguments.
+template <typename DominatorT>
+Block *findCommonDominator(Value value,
+                           const BufferViewFlowAnalysis::ValueSetT &values,
+                           const DominatorT &doms) {
+  // Store blocks in a set before querying `DominanceInfo` to filter out
+  // duplicate blocks (for performance reasons).
+  llvm::SmallPtrSet<Block *, 16> blocks;
+  // Start with the current block the value is defined in.
+  blocks.insert(value.getParentBlock());
+  for (Value childValue : values) {
+    for (Operation *user : childValue.getUsers()) {
+      // Find an appropriate dominator block that takes the current use into
+      // account.
+      blocks.insert(user->getBlock());
+    }
+    // Take values without any users into account.
+    blocks.insert(childValue.getParentBlock());
+  }
+  return doms.findNearestCommonDominator(blocks);
+}
+
 /// The base class for all BufferPlacement transformations.
 class BufferPlacementTransformationBase {
 public:
   using ValueSetT = BufferViewFlowAnalysis::ValueSetT;
 
-  /// Finds a common dominator for the given value while taking the positions
-  /// of the values in the value set into account. It supports dominator and
-  /// post-dominator analyses via template arguments.
-  template <typename DominatorT>
-  static Block *findCommonDominator(Value value, const ValueSetT &values,
-                                    const DominatorT &doms) {
-    // Start with the current block the value is defined in.
-    Block *dom = value.getParentBlock();
-    // Iterate over all aliases and their uses to find a safe placement block
-    // according to the given dominator information.
-    for (Value childValue : values) {
-      for (Operation *user : childValue.getUsers()) {
-        // Move upwards in the dominator tree to find an appropriate
-        // dominator block that takes the current use into account.
-        dom = doms.findNearestCommonDominator(dom, user->getBlock());
-      }
-      // Take values without any users into account.
-      dom = doms.findNearestCommonDominator(dom, childValue.getParentBlock());
-    }
-    return dom;
-  }
-
   /// Returns true if the given operation represents a loop by testing whether
   /// it implements the `LoopLikeOpInterface` or the `RegionBranchOpInterface`.
   /// In the case of a `RegionBranchOpInterface`, it checks all region-based
diff --git a/mlir/include/mlir/IR/Dominance.h b/mlir/include/mlir/IR/Dominance.h
index 82bf34c1910e86..2536ce585b3fdd 100644
--- a/mlir/include/mlir/IR/Dominance.h
+++ b/mlir/include/mlir/IR/Dominance.h
@@ -54,6 +54,21 @@ class DominanceInfoBase {
   /// nullptr.
   Block *findNearestCommonDominator(Block *a, Block *b) const;
 
+  /// Finds the nearest common dominator block for the given range of blocks.
+  /// If no common dominator can be found, this function will return nullptr.
+  template <typename BlockRangeT>
+  Block *findNearestCommonDominator(BlockRangeT &&blocks) const {
+    if (blocks.begin() == blocks.end())
+      return nullptr;
+    Block *dom = *blocks.begin();
+    for (auto it = ++blocks.begin(); it != blocks.end(); ++it) {
+      dom = findNearestCommonDominator(dom, *it);
+      if (!dom)
+        return nullptr;
+    }
+    return dom;
+  }
+
   /// Get the root dominance node of the given region. Note that this operation
   /// is only defined for multi-block regions!
   DominanceInfoNode *getRootNode(Region *region) {