[mlir][TosaToLinalg] Ensure valid insertion pointer after op erasure

[matthias-springer]

Erasing/replacing an op, which is also the current insertion point, invalidates the insertion point. Explicitly set the insertion point, so that subsequent op insertions do not crash after the One-Shot Dialect Conversion refactoring. (ConversionPatternRewriter will start behaving more like a "normal" rewriter.)

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-linalg

Author: Matthias Springer (matthias-springer)

Changes

Erasing/replacing an op, which is also the current insertion point, invalidates the insertion point. Explicitly set the insertion point, so that subsequent op insertions do not crash after the One-Shot Dialect Conversion refactoring. (ConversionPatternRewriter will start behaving more like a "normal" rewriter.)

---

Full diff: https://github.com/llvm/llvm-project/pull/146908.diff

1 Files Affected:

• (modified) mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp (+1)

diff --git a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
index c460a8bb2f4b2..ca8f6fcb1e122 100644
--- a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
+++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
@@ -815,6 +815,7 @@ class MaxPool2dConverter : public OpConversionPattern<tosa::MaxPool2dOp> {
         ValueRange{paddedInput, fakeWindowDims}, filledEmptyTensor, strideAttr,
         dilationAttr);
 
+    rewriter.setInsertionPointAfter(op);
     rewriter.replaceOp(op, resultOp);
 
     // NaN propagation has no meaning for non floating point types.

[llvmbot]

@llvm/pr-subscribers-mlir-tosa

Author: Matthias Springer (matthias-springer)

Changes

Erasing/replacing an op, which is also the current insertion point, invalidates the insertion point. Explicitly set the insertion point, so that subsequent op insertions do not crash after the One-Shot Dialect Conversion refactoring. (ConversionPatternRewriter will start behaving more like a "normal" rewriter.)

---

Full diff: https://github.com/llvm/llvm-project/pull/146908.diff

1 Files Affected:

• (modified) mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp (+1)

diff --git a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
index c460a8bb2f4b2..ca8f6fcb1e122 100644
--- a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
+++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
@@ -815,6 +815,7 @@ class MaxPool2dConverter : public OpConversionPattern<tosa::MaxPool2dOp> {
         ValueRange{paddedInput, fakeWindowDims}, filledEmptyTensor, strideAttr,
         dilationAttr);
 
+    rewriter.setInsertionPointAfter(op);
     rewriter.replaceOp(op, resultOp);
 
     // NaN propagation has no meaning for non floating point types.

[zero9178]

I am curious, do you think it'd be feasible for erase op (and replace op) to accommodate for the insertion point being erased? Eg move the insertion point to the next op if the current is being removed.

[matthias-springer]

I tried this here but it makes op erasure more expensive. What do you think?

[banach-space]

Thanks for the fix!

I tried this #146955 but it makes op erasure more expensive. What do you think?

Personally I feel that updating the insertion point should remain the responsibility of the user. Looking at this specific example, it feels like ewriter.replaceOp(op, resultOp); is invoked pre-maturely.

This is just an opinion, I don't have any data for or against a particular approach.