[mlir][sparse] fix memory leak sparse2sparse reshape

Reviewed By: aartbik

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

diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp
index c873600..d6f5bde 100644 (file)
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp
@@ -372,7 +372,9 @@ public:
           builder.create<sparse_tensor::YieldOp>(loc, t);
         });
     auto t = rewriter.create<LoadOp>(loc, foreachOp.getResult(0), true);
-    rewriter.replaceOpWithNewOp<ConvertOp>(op, dstTp, t);
+    auto converted = rewriter.create<ConvertOp>(loc, dstTp, t).getResult();
+    rewriter.create<DeallocTensorOp>(loc, t);
+    rewriter.replaceOp(op, converted);
     return success();
   }
 };