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

[llvmlistbot at llvm.org]

Author: Aviad Cohen
Date: 2024-04-07T08:23:16+03:00
New Revision: ccc02563f4d620d4d29a1cbd2c463871cc54745b

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

LOG: [mlir][linalg]: Fixed possible memory leak in cloneToCollapsedOp (#87595)

* Direct call to `clone` function leads to memory leak. Instead, we should use `RewriterBase` clone function instead.

Added: 
    

Modified: 
    mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
    mlir/lib/Dialect/Utils/StructuredOpsUtils.cpp

Removed: 
    


################################################################################
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);
 }