[Mlir-commits] [mlir] 8d0a02c - [mlir] Add InsertionGuards to OneToNPatternRewriter.

[Ingo Müller]

Author: Ingo Müller
Date: 2023-07-07T09:16:02Z
New Revision: 8d0a02cbc89eb9016062ecfa5abe88be45570cbb

URL: https://github.com/llvm/llvm-project/commit/8d0a02cbc89eb9016062ecfa5abe88be45570cbb
DIFF: https://github.com/llvm/llvm-project/commit/8d0a02cbc89eb9016062ecfa5abe88be45570cbb.diff

LOG: [mlir] Add InsertionGuards to OneToNPatternRewriter.

This fixes bad behavior of that class that surfaced in
https://reviews.llvm.org/D154299, where calling applySignatureConversion
left the insertion point different from before the call, which broke a
subsequent call to replaceOp. This patch introduces a fix in both
functions, each of which is enough to fix the specific problem in the
aforementioned diff: (1) applySignatureConversion now resets the
insertion point with a guard for the whole function and (2) replace sets
the insertion point to the op that should be replaced (and resets it
with a guard).

Reviewed By: ftynse

Differential Revision: https://reviews.llvm.org/D154684

Added: 
    

Modified: 
    mlir/lib/Transforms/Utils/OneToNTypeConversion.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Transforms/Utils/OneToNTypeConversion.cpp b/mlir/lib/Transforms/Utils/OneToNTypeConversion.cpp
index 44908510205f73..8c4df4bf8fe4d6 100644
--- a/mlir/lib/Transforms/Utils/OneToNTypeConversion.cpp
+++ b/mlir/lib/Transforms/Utils/OneToNTypeConversion.cpp
@@ -211,6 +211,8 @@ void OneToNPatternRewriter::replaceOp(Operation *op, ValueRange newValues,
   // original op with those.
   assert(newValues.size() == resultMapping.getConvertedTypes().size());
   assert(op->getResultTypes() == resultMapping.getOriginalTypes());
+  PatternRewriter::InsertionGuard g(*this);
+  setInsertionPointAfter(op);
   SmallVector<Value> castResults =
       buildUnrealizedBackwardsCasts(newValues, resultMapping, *this);
   replaceOp(op, castResults);
@@ -218,6 +220,8 @@ void OneToNPatternRewriter::replaceOp(Operation *op, ValueRange newValues,
 
 Block *OneToNPatternRewriter::applySignatureConversion(
     Block *block, OneToNTypeMapping &argumentConversion) {
+  PatternRewriter::InsertionGuard g(*this);
+
   // Split the block at the beginning to get a new block to use for the
   // updated signature.
   SmallVector<Location> locs;