From 61e5c14fa8d157c59ca25246f601be4e1fcba35c Mon Sep 17 00:00:00 2001
From: Peiming Liu <peiming@google.com>
Date: Wed, 16 Nov 2022 00:00:33 +0000
Subject: [PATCH] [mlir][sparse] fix memory leakage in concatenate rewriter.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D138074
---
 mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp
index 6fab361..2d2f903 100644
--- 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();
   }
 };
-- 
2.7.4