[mlir][linalg]: Fixed possible memory leak in cloneToCollapsedOp

[AviadCo]

• Direct call to clone function leads to memory leak. Instead, we would can use RewriterBase clone function.

[llvmbot]

@llvm/pr-subscribers-mlir

Author: Aviad Cohen (AviadCo)

Changes

• Direct call to clone function leads to memory leak. Instead, we would can use RewriterBase clone function.

---

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

1 Files Affected:

• (modified) mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp (+6-3)

diff --git a/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp b/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
index 9453502a253f16..34a5adec07bbea 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
@@ -1497,9 +1497,12 @@ LinalgOp cloneToCollapsedOp<LinalgOp>(RewriterBase &rewriter, LinalgOp origOp,
   SmallVector<Type> resultTypes;
   collapseOperandsAndResults(origOp, collapsingInfo, rewriter, inputOperands,
                              outputOperands, resultTypes);
-  return cast<LinalgOp>(clone(
-      rewriter, origOp, resultTypes,
-      llvm::to_vector(llvm::concat<Value>(inputOperands, outputOperands))));
+  IRMapping mapper;
+  for (auto [orig, input] : llvm::zip(origOp.getDpsInputs(), inputOperands))
+    mapper.map(orig, input);
+  for (auto [orig, output] : llvm::zip(origOp.getDpsInits(), outputOperands))
+    mapper.map(orig, output);
+  return cast<LinalgOp>(rewriter.clone(*origOp.getOperation(), mapper));
 }
 
 /// Collapse a `GenericOp`

[llvmbot]

@llvm/pr-subscribers-mlir-linalg

Author: Aviad Cohen (AviadCo)

Changes

• Direct call to clone function leads to memory leak. Instead, we would can use RewriterBase clone function.

---

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

1 Files Affected:

• (modified) mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp (+6-3)

diff --git a/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp b/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
index 9453502a253f16..34a5adec07bbea 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
@@ -1497,9 +1497,12 @@ LinalgOp cloneToCollapsedOp<LinalgOp>(RewriterBase &rewriter, LinalgOp origOp,
   SmallVector<Type> resultTypes;
   collapseOperandsAndResults(origOp, collapsingInfo, rewriter, inputOperands,
                              outputOperands, resultTypes);
-  return cast<LinalgOp>(clone(
-      rewriter, origOp, resultTypes,
-      llvm::to_vector(llvm::concat<Value>(inputOperands, outputOperands))));
+  IRMapping mapper;
+  for (auto [orig, input] : llvm::zip(origOp.getDpsInputs(), inputOperands))
+    mapper.map(orig, input);
+  for (auto [orig, output] : llvm::zip(origOp.getDpsInits(), outputOperands))
+    mapper.map(orig, output);
+  return cast<LinalgOp>(rewriter.clone(*origOp.getOperation(), mapper));
 }
 
 /// Collapse a `GenericOp`

[AviadCo]

@srcarroll I found this memory leak using Asan checks. Can you please take a look on the fix?

[MaheshRavishankar]

Thanks!

[srcarroll]

Oh interesting. Good catch. Thanks for the fix!

[AviadCo]

@MaheshRavishankar @srcarroll
I missed the right place of the memory leak, it actaully was due to direct call Region::cloneInto instead of using BuilderOp::cloneRegionBefore that also handle the listener, pushed new patch set for review.

[srcarroll]

since I dont understand why the original code causes memory leaks and why this fixes it, you should have @MaheshRavishankar or @nicolasvasilache give the final approval.