[Mlir-commits] [mlir] [mlir][sparse] deallocate temporary transposed tensor (PR #85720)

[Aart Bik]

https://github.com/aartbik created https://github.com/llvm/llvm-project/pull/85720

Last resort resolution of cycles introduced a sparse conversion without explicit sparse deallocation (which is not inserted by any automatic means). This fixes 2 out of 5 remaining asan detected leaks in sparse integration tests.

>From ef160ef5bc988b4aa3c2b4b43c32f0a9811b6028 Mon Sep 17 00:00:00 2001
From: Aart Bik <ajcbik at google.com>
Date: Mon, 18 Mar 2024 17:02:10 -0700
Subject: [PATCH] [mlir][sparse] deallocate temporary transposed tensor

Last resort resolution of cycles introduced a sparse conversion
without explicit sparse deallocation (which is not inserted by
any automatic means). This fixes 2 out of 5 remaining asan
detected leaks in sparse integration tests.
---
 .../SparseTensor/Transforms/SparseReinterpretMap.cpp        | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseReinterpretMap.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseReinterpretMap.cpp
index f93b59de29e57b..14ea07f0b54b82 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseReinterpretMap.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseReinterpretMap.cpp
@@ -573,6 +573,12 @@ struct GenericOpScheduler : public OpRewritePattern<linalg::GenericOp> {
       rewriter.modifyOpInPlace(linalgOp, [&]() {
         linalgOp->setOperand(t->getOperandNumber(), dst);
       });
+
+      // Release the transposed form afterwards.
+      // TODO: CSE when used in more than one following op?
+      rewriter.setInsertionPointAfter(linalgOp);
+      rewriter.create<bufferization::DeallocTensorOp>(dst.getLoc(), dst);
+
       return success();
     }
     // Cannot be resolved with a single conversion.