[Mlir-commits] [mlir] 3d90c81 - [mlir][sparse] Fix incorrect API usage in RewritePatterns

[Matthias Springer]

Author: Matthias Springer
Date: 2023-05-25T08:52:47+02:00
New Revision: 3d90c8126f8990490cd25a8a4b67a16b99898970

URL: https://github.com/llvm/llvm-project/commit/3d90c8126f8990490cd25a8a4b67a16b99898970
DIFF: https://github.com/llvm/llvm-project/commit/3d90c8126f8990490cd25a8a4b67a16b99898970.diff

LOG: [mlir][sparse] Fix incorrect API usage in RewritePatterns

Incorrect API usage was detected by
`MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS`.

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

Added: 
    

Modified: 
    mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp
index 6ee1c1b3dc49..ce9490cfddef 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp
@@ -367,7 +367,9 @@ struct FuseTensorCast : public OpRewritePattern<tensor::CastOp> {
     if (tensor::isSameTypeWithoutEncoding(srcType, dstType)) {
       if (Operation *def = op.getSource().getDefiningOp()) {
         if (def->hasOneUse() && isa<tensor::ExtractSliceOp>(def)) {
-          def->getResult(0).setType(op->getResultTypes()[0]);
+          rewriter.updateRootInPlace(def, [&]() {
+            def->getResult(0).setType(op->getResultTypes()[0]);
+          });
           rewriter.replaceOp(op, def->getResult(0));
           return success();
         }