[mlir][IR] Add helper functions to compute insertion point

Author: matthias-springer

mlir/include/mlir/IR/Builders.h

@@ -334,6 +334,18 @@ class OpBuilder : public Builder {
   /// This class represents a saved insertion point.
   class InsertPoint {
   public:
+    /// Finds the closest insertion point where all given values are defined
+    /// and can be used without a dominance violation. (Does not account for
+    /// IsolatedFromAbove regions.) Returns an "empty" insertion point if there
+    /// is no such insertion point.
+    static InsertPoint findClosest(ArrayRef<Value> values);
+
+    /// Creates a new insertion point right after the given value. If the given
+    /// value is a block argument, the insertion point points to the beginning
+    /// of the block. Otherwise, it points to the place right after the op that
+    /// defines the value.
+    static InsertPoint after(Value value);
+
     /// Creates a new insertion point which doesn't point to anything.
     InsertPoint() = default;
 

mlir/lib/IR/Builders.cpp

@@ -641,3 +641,76 @@ void OpBuilder::cloneRegionBefore(Region &region, Region &parent,
 void OpBuilder::cloneRegionBefore(Region &region, Block *before) {
   cloneRegionBefore(region, *before->getParent(), before->getIterator());
 }
+
+//===----------------------------------------------------------------------===//
+// InsertPoint
+//===----------------------------------------------------------------------===//
+
+OpBuilder::InsertPoint OpBuilder::InsertPoint::after(Value value) {
+  if (auto blockArg = dyn_cast<BlockArgument>(value))
+    return InsertPoint(blockArg.getOwner(), blockArg.getOwner()->begin());
+  Operation *op = value.getDefiningOp();
+  return InsertPoint(op->getBlock(), ++op->getIterator());
+}
+
+/// Helper function that returns "true" if:
+/// - `a` is a proper ancestor of `b`
+/// - or: there is a path from `a` to `b`
+static bool isAncestorOrBefore(Block *a, Block *b) {
+  if (a->getParentOp()->isProperAncestor(b->getParentOp()))
+    return true;
+  if (a->getParent() != b->getParent())
+    return false;
+  return a->isReachable(b);
+}
+
+OpBuilder::InsertPoint
+OpBuilder::InsertPoint::findClosest(ArrayRef<Value> values) {
+  // Compute the insertion point after the first value.
+  assert(!values.empty() && "expected at least one value");
+  InsertPoint result = InsertPoint::after(values.front());
+
+  // Check all other values and update the insertion point as needed.
+  for (Value v : values.drop_front()) {
+    InsertPoint pt = InsertPoint::after(v);
+
+    if (pt.getBlock() == result.getBlock()) {
+      // Both values belong to the same block. Modify the iterator (but keep
+      // the block) if needed: take the later one of the two insertion points.
+      Block *block = pt.getBlock();
+      if (pt.getPoint() == block->end()) {
+        // `pt` points to the end of the block: take `pt`.
+        result = pt;
+        continue;
+      } else if (result.getPoint() == block->end()) {
+        // `result` points to the end of the block: nothing to do.
+        continue;
+      }
+      // Neither `pt` nor `result` point to the end of the block, so both
+      // iterators point to an operation. Set `result` to the later one of the
+      // two insertion point.
+      if (result.getPoint()->isBeforeInBlock(&*pt.getPoint()))
+        result = pt;
+      continue;
+    }
+
+    if (isAncestorOrBefore(result.getBlock(), pt.getBlock())) {
+      // `result` is an ancestor of `pt`. Therefore, `pt` is a valid insertion
+      // point for `v` and all previous values.
+      result = pt;
+      continue;
+    }
+
+    if (isAncestorOrBefore(pt.getBlock(), result.getBlock())) {
+      // `pt` is an ancestor of `result`. Therefore, `result` is a valid
+      // insertion point for `v` and all previous values.
+      continue;
+    }
+
+    // `pt` and `result` are in different subtrees: neither is an ancestor of
+    // the other. In that case, there is no suitable insertion point.
+    return InsertPoint();
+  }
+
+  return result;
+}

mlir/lib/Transforms/Utils/DialectConversion.cpp

@@ -53,16 +53,6 @@ static void logFailure(llvm::ScopedPrinter &os, StringRef fmt, Args &&...args) {
   });
 }
 
-/// 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) {
-  Block *insertBlock = value.getParentBlock();
-  Block::iterator insertPt = insertBlock->begin();
-  if (OpResult inputRes = dyn_cast<OpResult>(value))
-    insertPt = ++inputRes.getOwner()->getIterator();
-  return OpBuilder::InsertPoint(insertBlock, insertPt);
-}
-
 //===----------------------------------------------------------------------===//
 // ConversionValueMapping
 //===----------------------------------------------------------------------===//
@@ -1147,8 +1137,9 @@ LogicalResult ConversionPatternRewriterImpl::remapValues(
       // that the value was replaced with a value of different type and no
       // source materialization was created yet.
       Value castValue = buildUnresolvedMaterialization(
-          MaterializationKind::Target, computeInsertPoint(newOperand),
-          operandLoc, /*inputs=*/newOperand, /*outputType=*/desiredType,
+          MaterializationKind::Target,
+          OpBuilder::InsertPoint::after(newOperand), operandLoc,
+          /*inputs=*/newOperand, /*outputType=*/desiredType,
           /*originalType=*/origType, currentTypeConverter);
       mapping.map(newOperand, castValue);
       newOperand = castValue;
@@ -1309,7 +1300,7 @@ Block *ConversionPatternRewriterImpl::applySignatureConversion(
     }
     if (legalOutputType && legalOutputType != origArgType) {
       Value targetMat = buildUnresolvedMaterialization(
-          MaterializationKind::Target, computeInsertPoint(argMat),
+          MaterializationKind::Target, OpBuilder::InsertPoint::after(argMat),
           origArg.getLoc(), /*inputs=*/argMat, /*outputType=*/legalOutputType,
           /*originalType=*/origArgType, converter);
       mapping.map(argMat, targetMat);
@@ -1401,7 +1392,7 @@ void ConversionPatternRewriterImpl::notifyOpReplaced(Operation *op,
 
       // Materialize a replacement value "out of thin air".
       newValue = buildUnresolvedMaterialization(
-          MaterializationKind::Source, computeInsertPoint(result),
+          MaterializationKind::Source, OpBuilder::InsertPoint::after(result),
           result.getLoc(), /*inputs=*/ValueRange(),
           /*outputType=*/result.getType(), /*originalType=*/Type(),
           currentTypeConverter);
@@ -2596,7 +2587,7 @@ void OperationConverter::finalize(ConversionPatternRewriter &rewriter) {
       Value newValue = rewriterImpl.mapping.lookupOrNull(originalValue);
       assert(newValue && "replacement value not found");
       Value castValue = rewriterImpl.buildUnresolvedMaterialization(
-          MaterializationKind::Source, computeInsertPoint(newValue),
+          MaterializationKind::Source, OpBuilder::InsertPoint::after(newValue),
           originalValue.getLoc(),
           /*inputs=*/newValue, /*outputType=*/originalValue.getType(),
           /*originalType=*/Type(), converter);