[Mlir-commits] [mlir] 61e5c14 - [mlir][sparse] fix memory leakage in concatenate rewriter.

[Peiming Liu]

Author: Peiming Liu
Date: 2022-11-16T00:10:33Z
New Revision: 61e5c14fa8d157c59ca25246f601be4e1fcba35c

URL: https://github.com/llvm/llvm-project/commit/61e5c14fa8d157c59ca25246f601be4e1fcba35c
DIFF: https://github.com/llvm/llvm-project/commit/61e5c14fa8d157c59ca25246f601be4e1fcba35c.diff

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

Reviewed By: aartbik

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

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 6fab36186d15..2d2f9031f6f1 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();
   }
 };