[mlir][sparse] deallocate temporary transposed tensor

[aartbik]

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.

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-sparse

Author: Aart Bik (aartbik)

Changes

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.

---

Full diff: https://github.com/llvm/llvm-project/pull/85720.diff

1 Files Affected:

• (modified) mlir/lib/Dialect/SparseTensor/Transforms/SparseReinterpretMap.cpp (+6)

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.