[mlir] Fix unsafe create operation in GreedyPatternRewriter

[mravishankar]

Summary

When trying to fold an operation during operation creation check that
the operation folding succeeds before inserting the op.

[rriddle]

mlir/include/mlir/Transforms/FoldUtils.h

78–87

Can you add comments as to why we don't use builder.create directly?

83

You will also want to insert in the case of successful fold and getNumResults() == 0. In that case it is treated as an in-place fold of a zero-result operation.

[mehdi_amini]

mlir/include/mlir/Transforms/FoldUtils.h

91

I think that is similar to what River mentions above, but this code is fishy: if none of the branch is taken the op just leaks.

I'd write it with early exit so that there is a clear handling of the "fallback" path.

[mravishankar]

Addressing comments

[mravishankar]

mlir/include/mlir/Transforms/FoldUtils.h

91

I changed it. I think the logic now maps to what River suggested. Something is off though. Test failing (maybe another issue exposed by this). Looking into those.

[mravishankar]

Fixing tests

[rriddle]

mlir/lib/Transforms/Utils/FoldUtils.cpp

169–170

Hmm, this is wrong. You can't assume the insertion point is on a valid operation, the block could be empty. You would need to use the block.

[rriddle]

mlir/lib/Transforms/Utils/FoldUtils.cpp

169–170

Passing in the insertion block instead of a builder would clean a bit of this up.

[mravishankar]

Addressing comments

[mravishankar]

mlir/lib/Transforms/Utils/FoldUtils.cpp

169–170

Good point. CHanged.

mlir/include/mlir/Transforms/FoldUtils.h

public:
   void notifyRemoval(Operation *op);
 
   /// Create an operation of specific op type with the given builder,
   /// and immediately try to fold it. This function populates 'results' with
   /// the results after folding the operation.
   template <typename OpTy, typename... Args>
   void create(OpBuilder &builder, SmallVectorImpl<Value> &results,
               Location location, Args &&... args) {
-    Operation *op = builder.create<OpTy>(location, std::forward<Args>(args)...);
-    if (failed(tryToFold(op, results)))
+    // The op needs to be inserted only if the fold (below) fails, or the number
+    // of results of the op is zero (which is treated as an in-place
+    // fold). Using create methods of the builder will insert the op, so not
+    // using it here.
+    OperationState state(location, OpTy::getOperationName());
+    OpTy::build(&builder, state, std::forward<Args>(args)...);
+    Operation *op = Operation::create(state);
+
+    if (failed(tryToFold(builder, op, results)) || op->getNumResults() == 0) {
+      builder.insert(op);
       results.assign(op->result_begin(), op->result_end());
-    else if (op->getNumResults() != 0)
-      op->erase();
+      return;
+    }
+    op->destroy();
   }
 
   /// Overload to create or fold a single result operation.
   template <typename OpTy, typename... Args>
   typename std::enable_if<OpTy::template hasTrait<OpTrait::OneResult>(),
                           Value>::type
   create(OpBuilder &builder, Location location, Args &&... args) {
     SmallVector<Value, 1> results;
private:
   /// generate different constants with the same input attribute and type, so we
   /// also need to track per-dialect.
   using ConstantMap =
       DenseMap<std::tuple<Dialect *, Attribute, Type>, Operation *>;
 
   /// Tries to perform folding on the given `op`. If successful, populates
   /// `results` with the results of the folding.
   LogicalResult tryToFold(
-      Operation *op, SmallVectorImpl<Value> &results,
+      OpBuilder &builder, Operation *op, SmallVectorImpl<Value> &results,
       function_ref<void(Operation *)> processGeneratedConstants = nullptr);
 
   /// Try to get or create a new constant entry. On success this returns the
   /// constant operation, nullptr otherwise.
   Operation *tryGetOrCreateConstant(ConstantMap &uniquedConstants,
                                     Dialect *dialect, OpBuilder &builder,
                                     Attribute value, Type type, Location loc);
 

mlir/lib/Transforms/Utils/FoldUtils.cpp

 #include "mlir/IR/Operation.h"
 
 using namespace mlir;
 
 /// Given an operation, find the parent region that folded constants should be
 /// inserted into.
 static Region *getInsertionRegion(
     DialectInterfaceCollection<OpFolderDialectInterface> &interfaces,
-    Operation *op) {
-  while (Region *region = op->getParentRegion()) {
+    Block *insertionBlock) {
+  while (Region *region = insertionBlock->getParent()) {
     // Insert in this region for any of the following scenarios:
     //  * The parent is unregistered, or is known to be isolated from above.
     //  * The parent is a top-level operation.
     auto *parentOp = region->getParentOp();
     if (!parentOp->isRegistered() || parentOp->isKnownIsolatedFromAbove() ||
         !parentOp->getBlock())
       return region;
 
     // Otherwise, check if this region is a desired insertion region.
     auto *interface = interfaces.getInterfaceFor(parentOp);
     if (LLVM_UNLIKELY(interface && interface->shouldMaterializeInto(region)))
       return region;
 
     // Traverse up the parent looking for an insertion region.
-    op = parentOp;
+    insertionBlock = parentOp->getBlock();
   }
   llvm_unreachable("expected valid insertion region");
 }
 
 /// A utility function used to materialize a constant for a given attribute and
 /// type. On success, a valid constant value is returned. Otherwise, null is
 /// returned
 static Operation *materializeConstant(Dialect *dialect, OpBuilder &builder,
LogicalResult OperationFolder::tryToFold(
     function_ref<void(Operation *)> preReplaceAction) {
   // If this is a unique'd constant, return failure as we know that it has
   // already been folded.
   if (referencedDialects.count(op))
     return failure();
 
   // Try to fold the operation.
   SmallVector<Value, 8> results;
-  if (failed(tryToFold(op, results, processGeneratedConstants)))
+  OpBuilder builder(op);
+  if (failed(tryToFold(builder, op, results, processGeneratedConstants)))
     return failure();
 
   // Check to see if the operation was just updated in place.
   if (results.empty())
     return success();
 
   // Constant folding succeeded. We will start replacing this op's uses and
   // erase this op. Invoke the callback provided by the caller to perform any
void OperationFolder::notifyRemoval(Operation *op) {
 
   // Get the constant value for this operation, this is the value that was used
   // to unique the operation internally.
   Attribute constValue;
   matchPattern(op, m_Constant(&constValue));
   assert(constValue);
 
   // Get the constant map that this operation was uniqued in.
-  auto &uniquedConstants = foldScopes[getInsertionRegion(interfaces, op)];
+  auto &uniquedConstants =
+      foldScopes[getInsertionRegion(interfaces, op->getBlock())];
 
   // Erase all of the references to this operation.
   auto type = op->getResult(0).getType();
   for (auto *dialect : it->second)
     uniquedConstants.erase(std::make_tuple(dialect, constValue, type));
   referencedDialects.erase(it);
 }
 
 /// Clear out any constants cached inside of the folder.
 void OperationFolder::clear() {
   foldScopes.clear();
   referencedDialects.clear();
 }
 
 /// Tries to perform folding on the given `op`. If successful, populates
 /// `results` with the results of the folding.
 LogicalResult OperationFolder::tryToFold(
-    Operation *op, SmallVectorImpl<Value> &results,
+    OpBuilder &builder, Operation *op, SmallVectorImpl<Value> &results,
     function_ref<void(Operation *)> processGeneratedConstants) {
   SmallVector<Attribute, 8> operandConstants;
   SmallVector<OpFoldResult, 8> foldResults;
 
   // If this is a commutative operation, move constants to be trailing operands.
   if (op->getNumOperands() >= 2 && op->isCommutative()) {
     std::stable_partition(
         op->getOpOperands().begin(), op->getOpOperands().end(),
LogicalResult OperationFolder::tryToFold(
 
   // Check to see if the operation was just updated in place.
   if (foldResults.empty())
     return success();
   assert(foldResults.size() == op->getNumResults());
 
   // Create a builder to insert new operations into the entry block of the
   // insertion region.
-  auto *insertRegion = getInsertionRegion(interfaces, op);
+  auto *insertRegion =
+      getInsertionRegion(interfaces, builder.getInsertionBlock());
   auto &entry = insertRegion->front();
-  OpBuilder builder(&entry, entry.begin());
+  OpBuilder::InsertionGuard foldGuard(builder);
+  builder.setInsertionPoint(&entry, entry.begin());
 
   // Get the constant map for the insertion region of this operation.
   auto &uniquedConstants = foldScopes[insertRegion];
 
   // Create the result constants and replace the results.
   auto *dialect = op->getDialect();
   for (unsigned i = 0, e = op->getNumResults(); i != e; ++i) {
     assert(!foldResults[i].isNull() && "expected valid OpFoldResult");