[Mlir-commits] [mlir] 8ffdcc5 - [mlir][sparse] fix memory leak sparse2sparse reshape

[Peiming Liu]

Author: Peiming Liu
Date: 2022-11-15T00:19:51Z
New Revision: 8ffdcc594e8bd4807b5f0fd4f382523ef6a17ae4

URL: https://github.com/llvm/llvm-project/commit/8ffdcc594e8bd4807b5f0fd4f382523ef6a17ae4
DIFF: https://github.com/llvm/llvm-project/commit/8ffdcc594e8bd4807b5f0fd4f382523ef6a17ae4.diff

LOG: [mlir][sparse] fix memory leak sparse2sparse reshape

Reviewed By: aartbik

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

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 c873600f891b..d6f5bded2e86 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp
@@ -372,7 +372,9 @@ struct Sparse2SparseReshapeRewriter : public OpRewritePattern<ReshapeOp> {
           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();
   }
 };