[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

168 ↗

(On Diff #251415)

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

168 ↗

(On Diff #251415)

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

168 ↗

(On Diff #251415)

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(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;