[Mlir-commits] [mlir] 6958255 - [mlir][TOSA] Do not access erased operation (#148374)

[llvmlistbot at llvm.org]

Author: Matthias Springer
Date: 2025-07-14T10:43:00+02:00
New Revision: 695825589dc1e91272f1997ad92ab58cd8f7ca70

URL: https://github.com/llvm/llvm-project/commit/695825589dc1e91272f1997ad92ab58cd8f7ca70
DIFF: https://github.com/llvm/llvm-project/commit/695825589dc1e91272f1997ad92ab58cd8f7ca70.diff

LOG: [mlir][TOSA] Do not access erased operation (#148374)

This commit fixes two occurrences where an erased op was accessed at a
later point of time. That won't work anymore in a One-Shot Dialect
Conversion and triggers a use-after-free sanitizer error.

After the One-Shot Dialect Conversion refactoring, a
`ConversionPatternRewriter` will behave more like a normal
`PatternRewriter`.

---------

Co-authored-by: Markus Böck <markus.boeck02 at gmail.com>

Added: 
    

Modified: 
    mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
    mlir/lib/Conversion/TosaToTensor/TosaToTensor.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
index c1f40dcbd5ca0..00b9a065dfb3d 100644
--- a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
+++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
@@ -808,6 +808,7 @@ class MaxPool2dConverter : public OpConversionPattern<tosa::MaxPool2dOp> {
         dilationAttr);
 
     rewriter.setInsertionPointAfter(op);
+    StringRef nanMode = op.getNanMode();
     rewriter.replaceOp(op, resultOp);
 
     // NaN propagation has no meaning for non floating point types.
@@ -821,11 +822,10 @@ class MaxPool2dConverter : public OpConversionPattern<tosa::MaxPool2dOp> {
     // we've already produced a named op we will just take its body and modify
     // it to include the appropriate checks. If the current value is NaN the
     // old value of pool will be taken otherwise we use the result.
-    if (const auto nanMode = op.getNanMode(); nanMode == "IGNORE") {
+    if (nanMode == "IGNORE") {
       auto genericOp = rewriter.create<linalg::GenericOp>(
-          op->getLoc(), resultOp.getType(0), resultOp.getInputs(),
-          resultOp.getOutputs(), resultOp.getIndexingMapsArray(),
-          resultOp.getIteratorTypesArray(),
+          loc, resultOp.getType(0), resultOp.getInputs(), resultOp.getOutputs(),
+          resultOp.getIndexingMapsArray(), resultOp.getIteratorTypesArray(),
           [&](OpBuilder &opBuilder, Location loc, ValueRange blockArgs) {
             IRMapping map;
             auto oldBlock = resultOp.getRegion().begin();
@@ -834,10 +834,10 @@ class MaxPool2dConverter : public OpConversionPattern<tosa::MaxPool2dOp> {
             map.map(oldArgs, blockArgs);
             auto *newOp = opBuilder.clone(oldMaxOp, map);
             Value isNaN = opBuilder.create<arith::CmpFOp>(
-                op->getLoc(), arith::CmpFPredicate::UNO, blockArgs.front(),
+                loc, arith::CmpFPredicate::UNO, blockArgs.front(),
                 blockArgs.front());
             auto selectOp = opBuilder.create<arith::SelectOp>(
-                op->getLoc(), isNaN, blockArgs.back(), newOp->getResult(0));
+                loc, isNaN, blockArgs.back(), newOp->getResult(0));
             opBuilder.create<linalg::YieldOp>(loc, selectOp.getResult());
           });
       rewriter.replaceOp(resultOp, genericOp);

diff  --git a/mlir/lib/Conversion/TosaToTensor/TosaToTensor.cpp b/mlir/lib/Conversion/TosaToTensor/TosaToTensor.cpp
index 9ceb5c0c7f2fe..c6cbcb0f8ab2b 100644
--- a/mlir/lib/Conversion/TosaToTensor/TosaToTensor.cpp
+++ b/mlir/lib/Conversion/TosaToTensor/TosaToTensor.cpp
@@ -320,8 +320,6 @@ class SliceConverter : public OpConversionPattern<tosa::SliceOp> {
         rewriter.getDenseI64ArrayAttr(sizes),
         rewriter.getDenseI64ArrayAttr(strides));
 
-    rewriter.replaceOp(sliceOp, newSliceOp.getResult());
-
     // Remove const_shape ops when it no longer has use point.
     Operation *startConstShape = sliceOp.getStart().getDefiningOp();
     if (startConstShape->getResult(0).hasOneUse())
@@ -331,6 +329,7 @@ class SliceConverter : public OpConversionPattern<tosa::SliceOp> {
     if (sizeConstShape->getResult(0).hasOneUse())
       rewriter.eraseOp(sizeConstShape);
 
+    rewriter.replaceOp(sliceOp, newSliceOp.getResult());
     return success();
   }
 };