[mlir][Transforms][NFC] Use DominanceInfo to compute materialization insertion point

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 properlyDominates(Block *aBlock, Block::iterator aIt, Block *bBlock,
+                         Block::iterator bIt, bool enclosingOk = true) const {
+    return super::properlyDominatesImpl(aBlock, aIt, bBlock, bIt, enclosingOk);
+  }
+
+  bool dominates(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 properlyPostDominates(Block *aBlock, Block::iterator aIt, Block *bBlock,
+                             Block::iterator bIt,
+                             bool enclosingOk = true) const {
+    return super::properlyDominatesImpl(aBlock, aIt, bBlock, bIt, enclosingOk);
+  }
+
+  bool postDominates(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/Transforms/Utils/DialectConversion.cpp

@@ -54,55 +54,6 @@ static void logFailure(llvm::ScopedPrinter &os, StringRef fmt, Args &&...args) {
   });
 }
 
-/// Given two insertion points in the same block, choose the later one.
-static OpBuilder::InsertPoint
-chooseLaterInsertPointInBlock(OpBuilder::InsertPoint a,
-                              OpBuilder::InsertPoint b) {
-  assert(a.getBlock() == b.getBlock() && "expected same block");
-  Block *block = a.getBlock();
-  if (a.getPoint() == block->begin())
-    return b;
-  if (b.getPoint() == block->begin())
-    return a;
-  if (a.getPoint()->isBeforeInBlock(&*b.getPoint()))
-    return b;
-  return a;
-}
-
-/// Helper function that chooses the insertion point among the two given ones
-/// that is later.
-// TODO: Extend DominanceInfo API to work with block iterators.
-static OpBuilder::InsertPoint chooseLaterInsertPoint(OpBuilder::InsertPoint a,
-                                                     OpBuilder::InsertPoint b) {
-  // Case 1: Fast path: Same block. This is the most common case.
-  if (LLVM_LIKELY(a.getBlock() == b.getBlock()))
-    return chooseLaterInsertPointInBlock(a, b);
-
-  // Case 2: Different block, but same region.
-  if (a.getBlock()->getParent() == b.getBlock()->getParent()) {
-    DominanceInfo domInfo;
-    if (domInfo.properlyDominates(a.getBlock(), b.getBlock()))
-      return b;
-    if (domInfo.properlyDominates(b.getBlock(), a.getBlock()))
-      return a;
-    // Neither of the two blocks dominante each other.
-    llvm_unreachable("unable to find valid insertion point");
-  }
-
-  // Case 3: b's region contains a: choose a.
-  if (b.getBlock()->getParent()->findAncestorOpInRegion(
-          *a.getPoint()->getParentOp()))
-    return a;
-
-  // Case 4: a's region contains b: choose b.
-  if (a.getBlock()->getParent()->findAncestorOpInRegion(
-          *b.getPoint()->getParentOp()))
-    return b;
-
-  // Neither of the two operations contain each other.
-  llvm_unreachable("unable to find valid insertion point");
-}
-
 /// Helper function that computes an insertion point where the given value is
 /// defined and can be used without a dominance violation.
 static OpBuilder::InsertPoint computeInsertPoint(Value value) {
@@ -117,9 +68,26 @@ static OpBuilder::InsertPoint computeInsertPoint(Value value) {
 /// defined and can be used without a dominance violation.
 static OpBuilder::InsertPoint computeInsertPoint(ArrayRef<Value> vals) {
   assert(!vals.empty() && "expected at least one value");
+  DominanceInfo domInfo;
   OpBuilder::InsertPoint pt = computeInsertPoint(vals.front());
-  for (Value v : vals.drop_front())
-    pt = chooseLaterInsertPoint(pt, computeInsertPoint(v));
+  for (Value v : vals.drop_front()) {
+    // Choose the "later" insertion point.
+    OpBuilder::InsertPoint nextPt = computeInsertPoint(v);
+    if (domInfo.dominates(pt.getBlock(), pt.getPoint(), nextPt.getBlock(),
+                          nextPt.getPoint())) {
+      // pt is before nextPt => choose nextPt.
+      pt = nextPt;
+    } else {
+#ifndef NDEBUG
+      // nextPt should be before pt => choose pt.
+      // If pt, nextPt are no dominance relationship, then there is no valid
+      // insertion point at which all given values are defined.
+      bool dom = domInfo.dominates(nextPt.getBlock(), nextPt.getPoint(),
+                                   pt.getBlock(), pt.getPoint());
+      assert(dom && "unable to find valid insertion point");
+#endif // NDEBUG
+    }
+  }
   return pt;
 }