[llvm-branch-commits] [mlir] [mlir][IR] Add `OpBuilder::setInsertionPointAfterValues` (PR #114940)

[Matthias Springer via llvm-branch-commits]

https://github.com/matthias-springer updated https://github.com/llvm/llvm-project/pull/114940

>From f716eaa7bfa4bf2e3971a908009255cccfd631df Mon Sep 17 00:00:00 2001
From: Matthias Springer <mspringer at nvidia.com>
Date: Sat, 9 Nov 2024 12:29:16 +0100
Subject: [PATCH] [mlir][IR] Add `OpBuilder::setInsertionPointAfterValues`

---
 mlir/include/mlir/IR/Builders.h  | 14 +++++++++++
 mlir/include/mlir/IR/Dominance.h | 23 +++++++++++++++++
 mlir/lib/IR/Builders.cpp         | 42 ++++++++++++++++++++++++++++++++
 3 files changed, 79 insertions(+)

diff --git a/mlir/include/mlir/IR/Builders.h b/mlir/include/mlir/IR/Builders.h
index 6fb71ccefda151..bd3642f9f413dc 100644
--- a/mlir/include/mlir/IR/Builders.h
+++ b/mlir/include/mlir/IR/Builders.h
@@ -16,6 +16,7 @@
 namespace mlir {
 
 class AffineExpr;
+class PostDominanceInfo;
 class IRMapping;
 class UnknownLoc;
 class FileLineColLoc;
@@ -435,6 +436,19 @@ class OpBuilder : public Builder {
     }
   }
 
+  /// Sets the insertion point to a place that post-dominates the definitions
+  /// of all given values. Returns "failure" and leaves the current insertion
+  /// point unchanged if no such insertion point exists.
+  ///
+  /// There may be multiple suitable insertion points. This function chooses an
+  /// insertion right after one of the given values.
+  ///
+  /// Note: Some of the given values may already have gone out of scope at the
+  /// selected insertion point. (E.g., because they are defined in a nested
+  /// region.)
+  LogicalResult setInsertionPointAfterValues(ArrayRef<Value> values,
+                                             const PostDominanceInfo &domInfo);
+
   /// Sets the insertion point to the start of the specified block.
   void setInsertionPointToStart(Block *block) {
     setInsertionPoint(block, block->begin());
diff --git a/mlir/include/mlir/IR/Dominance.h b/mlir/include/mlir/IR/Dominance.h
index 63504cad211a4d..be2dcec380b6cc 100644
--- a/mlir/include/mlir/IR/Dominance.h
+++ b/mlir/include/mlir/IR/Dominance.h
@@ -187,6 +187,17 @@ class DominanceInfo : public detail::DominanceInfoBase</*IsPostDom=*/false> {
   /// dominance" of ops, the single block is considered to properly dominate
   /// itself in a graph region.
   bool properlyDominates(Block *a, Block *b) const;
+
+  bool properlyDominantes(Block *aBlock, Block::iterator aIt, Block *bBlock,
+                          Block::iterator bIt, bool enclosingOk = true) const {
+    return super::properlyDominatesImpl(aBlock, aIt, bBlock, bIt, enclosingOk);
+  }
+
+  bool dominantes(Block *aBlock, Block::iterator aIt, Block *bBlock,
+                  Block::iterator bIt, bool enclosingOk = true) const {
+    return (aBlock == bBlock && aIt == bIt) ||
+           super::properlyDominatesImpl(aBlock, aIt, bBlock, bIt, enclosingOk);
+  }
 };
 
 /// A class for computing basic postdominance information.
@@ -210,6 +221,18 @@ class PostDominanceInfo : public detail::DominanceInfoBase</*IsPostDom=*/true> {
   bool postDominates(Block *a, Block *b) const {
     return a == b || properlyPostDominates(a, b);
   }
+
+  bool properlyPostDominantes(Block *aBlock, Block::iterator aIt, Block *bBlock,
+                              Block::iterator bIt,
+                              bool enclosingOk = true) const {
+    return super::properlyDominatesImpl(aBlock, aIt, bBlock, bIt, enclosingOk);
+  }
+
+  bool postDominantes(Block *aBlock, Block::iterator aIt, Block *bBlock,
+                      Block::iterator bIt, bool enclosingOk = true) const {
+    return (aBlock == bBlock && aIt == bIt) ||
+           super::properlyDominatesImpl(aBlock, aIt, bBlock, bIt, enclosingOk);
+  }
 };
 
 } // namespace mlir
diff --git a/mlir/lib/IR/Builders.cpp b/mlir/lib/IR/Builders.cpp
index 5397fbabc5c95e..0f4b5123326502 100644
--- a/mlir/lib/IR/Builders.cpp
+++ b/mlir/lib/IR/Builders.cpp
@@ -11,6 +11,7 @@
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Dialect.h"
+#include "mlir/IR/Dominance.h"
 #include "mlir/IR/IRMapping.h"
 #include "mlir/IR/IntegerSet.h"
 #include "mlir/IR/Matchers.h"
@@ -641,3 +642,44 @@ void OpBuilder::cloneRegionBefore(Region &region, Region &parent,
 void OpBuilder::cloneRegionBefore(Region &region, Block *before) {
   cloneRegionBefore(region, *before->getParent(), before->getIterator());
 }
+
+LogicalResult
+OpBuilder::setInsertionPointAfterValues(ArrayRef<Value> values,
+                                        const PostDominanceInfo &domInfo) {
+  // Helper function that computes the point after v's definition.
+  auto computeAfterIp = [](Value v) -> std::pair<Block *, Block::iterator> {
+    if (auto blockArg = dyn_cast<BlockArgument>(v))
+      return std::make_pair(blockArg.getOwner(), blockArg.getOwner()->begin());
+    Operation *op = v.getDefiningOp();
+    return std::make_pair(op->getBlock(), op->getIterator());
+  };
+
+  // Compute the insertion point after the first value is defined.
+  assert(!values.empty() && "expected at least one Value");
+  auto [block, blockIt] = computeAfterIp(values.front());
+
+  // Check the other values one-by-one and update the insertion point if
+  // needed.
+  for (Value v : values.drop_front()) {
+    auto [candidateBlock, candidateBlockIt] = computeAfterIp(v);
+    if (domInfo.postDominantes(candidateBlock, candidateBlockIt, block,
+                               blockIt)) {
+      // The point after v's definition post-dominates the current (and all
+      // previous) insertion points. Note: Post-dominance is transitive.
+      block = candidateBlock;
+      blockIt = candidateBlockIt;
+      continue;
+    }
+
+    if (!domInfo.postDominantes(block, blockIt, candidateBlock,
+                                candidateBlockIt)) {
+      // The point after v's definition and the current insertion point do not
+      // post-dominate each other. Therefore, there is no insertion point that
+      // post-dominates all values.
+      return failure();
+    }
+  }
+
+  setInsertionPoint(block, blockIt);
+  return success();
+}