[Mlir-commits] [mlir] [mlir][IR] `DominanceInfo`: Add function to query dominator of a range of block (PR #77098)

[Matthias Springer]

https://github.com/matthias-springer created https://github.com/llvm/llvm-project/pull/77098

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

>From 96b91217bb41a9cf2f2b84de3ae2e4189f23e26c Mon Sep 17 00:00:00 2001
From: Matthias Springer <springerm at google.com>
Date: Fri, 5 Jan 2024 14:27:07 +0000
Subject: [PATCH] [mlir][IR] `DominanceInfo`: Add function to query dominator
 of a range of block

Also improve the implementation of `findCommonDominator` and extract it
from `BufferPlacementTransformationBase` (so that
`BufferPlacementTransformationBase` can be retired eventually).
---
 .../Bufferization/Transforms/BufferUtils.h    | 46 ++++++++++---------
 mlir/include/mlir/IR/Dominance.h              | 15 ++++++
 2 files changed, 39 insertions(+), 22 deletions(-)

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) {