[Mlir-commits] [mlir] [mlir][linalg]: Fixed possible memory leak in cloneToCollapsedOp (PR #87595)

[Aviad Cohen]

https://github.com/AviadCo updated https://github.com/llvm/llvm-project/pull/87595

>From d03112bdcd07a283d0610a709030c502062356a3 Mon Sep 17 00:00:00 2001
From: Aviad Cohen <aviad.cohen2 at mobileye.com>
Date: Thu, 4 Apr 2024 06:38:21 +0300
Subject: [PATCH] [mlir][linalg]: Fixed possible memory leak in
 cloneToCollapsedOp

* Direct call to `clone` function leads to memory leak. Instead, we
would can use `RewriterBase` clone function.
---
 mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp | 5 +++--
 mlir/lib/Dialect/Utils/StructuredOpsUtils.cpp              | 6 ++++--
 2 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp b/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
index 9453502a253f16..373e9cfc3ce719 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
@@ -1497,9 +1497,10 @@ LinalgOp cloneToCollapsedOp<LinalgOp>(RewriterBase &rewriter, LinalgOp origOp,
   SmallVector<Type> resultTypes;
   collapseOperandsAndResults(origOp, collapsingInfo, rewriter, inputOperands,
                              outputOperands, resultTypes);
-  return cast<LinalgOp>(clone(
+
+  return clone(
       rewriter, origOp, resultTypes,
-      llvm::to_vector(llvm::concat<Value>(inputOperands, outputOperands))));
+      llvm::to_vector(llvm::concat<Value>(inputOperands, outputOperands)));
 }
 
 /// Collapse a `GenericOp`
diff --git a/mlir/lib/Dialect/Utils/StructuredOpsUtils.cpp b/mlir/lib/Dialect/Utils/StructuredOpsUtils.cpp
index 383ef1cea53fd3..adde8a66d83547 100644
--- a/mlir/lib/Dialect/Utils/StructuredOpsUtils.cpp
+++ b/mlir/lib/Dialect/Utils/StructuredOpsUtils.cpp
@@ -199,8 +199,10 @@ Operation *mlir::clone(OpBuilder &b, Operation *op, TypeRange newResultTypes,
   IRMapping bvm;
   OperationState state(op->getLoc(), op->getName(), newOperands, newResultTypes,
                        op->getAttrs());
-  for (Region &r : op->getRegions())
-    r.cloneInto(state.addRegion(), bvm);
+  for (Region &r : op->getRegions()) {
+    Region *newRegion = state.addRegion();
+    b.cloneRegionBefore(r, *newRegion, newRegion->begin(), bvm);
+  }
   return b.create(state);
 }