[mlir][IR] Add InsertPoint::after(ValueRange)

mlir/include/mlir/IR/Builders.h

@@ -16,6 +16,7 @@
 namespace mlir {
 
 class AffineExpr;
+class PostDominanceInfo;
 class IRMapping;
 class UnknownLoc;
 class FileLineColLoc;
@@ -341,6 +342,19 @@ class OpBuilder : public Builder {
     InsertPoint(Block *insertBlock, Block::iterator insertPt)
         : block(insertBlock), point(insertPt) {}
 
+    /// Compute an insertion point to a place that post-dominates the
+    /// definitions of all given values. Returns an "empty" insertion point 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 or because they are not visible in an IsolatedFromAbove region.)
+    static InsertPoint after(ArrayRef<Value> values,
+                             const PostDominanceInfo &domInfo);
+
     /// Returns true if this insert point is set.
     bool isSet() const { return (block != nullptr); }
 

mlir/include/mlir/IR/Dominance.h

@@ -113,8 +113,12 @@ class DominanceInfoBase {
   llvm::PointerIntPair<DomTree *, 1, bool>
   getDominanceInfo(Region *region, bool needsDomTree) const;
 
-  /// Return true if the specified block A properly dominates block B.
-  bool properlyDominates(Block *a, Block *b) const;
+  /// Return "true" if block iterator A properly (post)dominates block iterator
+  /// B. If `enclosingOk` is set, A is considered to (post)dominate B if A
+  /// encloses B.
+  bool properlyDominatesImpl(Block *aBlock, Block::iterator aIt, Block *bBlock,
+                             Block::iterator bIt,
+                             bool enclosingOk = true) const;
 
   /// A mapping of regions to their base dominator tree and a cached
   /// "hasSSADominance" bit. This map does not contain dominator trees for
@@ -147,9 +151,7 @@ class DominanceInfo : public detail::DominanceInfoBase</*IsPostDom=*/false> {
   /// The `enclosingOpOk` flag says whether we should return true if the B op
   /// is enclosed by a region on A.
   bool properlyDominates(Operation *a, Operation *b,
-                         bool enclosingOpOk = true) const {
-    return properlyDominatesImpl(a, b, enclosingOpOk);
-  }
+                         bool enclosingOpOk = true) const;
 
   /// Return true if operation A dominates operation B, i.e. if A and B are the
   /// same operation or A properly dominates B.
@@ -184,16 +186,18 @@ class DominanceInfo : public detail::DominanceInfoBase</*IsPostDom=*/false> {
   /// Graph regions have only a single block. To be consistent with "proper
   /// dominance" of ops, the single block is considered to properly dominate
   /// itself in a graph region.
-  bool properlyDominates(Block *a, Block *b) const {
-    return super::properlyDominates(a, b);
+  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);
   }
 
-private:
-  // Return true if operation A properly dominates operation B.  The
-  /// 'enclosingOpOk' flag says whether we should return true if the b op is
-  /// enclosed by a region on 'A'.
-  bool properlyDominatesImpl(Operation *a, Operation *b,
-                             bool enclosingOpOk) const;
+  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.
@@ -202,22 +206,33 @@ class PostDominanceInfo : public detail::DominanceInfoBase</*IsPostDom=*/true> {
   using super::super;
 
   /// Return true if operation A properly postdominates operation B.
-  bool properlyPostDominates(Operation *a, Operation *b) const;
+  bool properlyPostDominates(Operation *a, Operation *b,
+                             bool enclosingOpOk = true) const;
 
   /// Return true if operation A postdominates operation B.
   bool postDominates(Operation *a, Operation *b) const {
     return a == b || properlyPostDominates(a, b);
   }
 
   /// Return true if the specified block A properly postdominates block B.
-  bool properlyPostDominates(Block *a, Block *b) const {
-    return super::properlyDominates(a, b);
-  }
+  bool properlyPostDominates(Block *a, Block *b) const;
 
   /// Return true if the specified block A postdominates block B.
   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

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,43 @@ void OpBuilder::cloneRegionBefore(Region &region, Region &parent,
 void OpBuilder::cloneRegionBefore(Region &region, Block *before) {
   cloneRegionBefore(region, *before->getParent(), before->getIterator());
 }
+
+OpBuilder::InsertPoint
+OpBuilder::InsertPoint::after(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 InsertPoint(nullptr, Block::iterator());
+    }
+  }
+
+  return InsertPoint(block, blockIt);
+}