[mlir][sparse] fix memory leakage in concatenate rewriter.

Reviewed By: aartbik

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

diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp
index 6fab361..2d2f903 100644 (file)
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp
@@ -499,7 +499,10 @@ struct ConcatenateRewriter : public OpRewritePattern<ConcatenateOp> {
     }
 
     cooBuffer = rewriter.create<LoadOp>(loc, cooBuffer, true);
-    rewriter.replaceOpWithNewOp<ConvertOp>(op, rtp, cooBuffer);
+    Value converted =
+        rewriter.create<ConvertOp>(loc, rtp, cooBuffer).getResult();
+    rewriter.create<DeallocTensorOp>(loc, cooBuffer);
+    rewriter.replaceOp(op, converted);
     return success();
   }
 };