From e7061c59c1cb5e37d1bf9dfb5c7b8529b40cd9f6 Mon Sep 17 00:00:00 2001
From: snigdha dalvi <sdalvi@quicinc.com>
Date: Wed, 10 Sep 2025 12:28:33 -0700
Subject: [PATCH 1/2] Add flag to guard MoveInitOperandsToInputPattern in
 LinalgFoldUnitExtentDimsPass, as it causes regression in performance.

---
 mlir/include/mlir/Dialect/Linalg/Passes.td          | 5 ++++-
 mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp | 4 +++-
 2 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/mlir/include/mlir/Dialect/Linalg/Passes.td b/mlir/include/mlir/Dialect/Linalg/Passes.td
index 44da2965e6892..7da60bef6d70c 100644
--- a/mlir/include/mlir/Dialect/Linalg/Passes.td
+++ b/mlir/include/mlir/Dialect/Linalg/Passes.td
@@ -145,7 +145,10 @@ def LinalgFoldUnitExtentDimsPass : Pass<"linalg-fold-unit-extent-dims", ""> {
   let options = [
     Option<"useRankReducingSlices", "use-rank-reducing-slices", "bool",
            /*default=*/"false",
-           "Generate rank-reducing slices instead of reassociative reshapes">
+           "Generate rank-reducing slices instead of reassociative reshapes">,
+    Option<"enableMoveInitOperandsToInput", "enable-move-init-operands-to-input", "bool",
+           /*default=*/"false",
+           "Enable MoveInitOperandsToInputPattern transformation">
   ];
   let dependentDialects = [
     "linalg::LinalgDialect", "affine::AffineDialect", "memref::MemRefDialect"
diff --git a/mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp b/mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp
index 22690daa4f9e1..ba466ed3df5cd 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp
@@ -868,7 +868,9 @@ struct LinalgFoldUnitExtentDimsPass
           RankReductionStrategy::ExtractInsertSlice;
     }
     linalg::populateFoldUnitExtentDimsPatterns(patterns, options);
-    populateMoveInitOperandsToInputPattern(patterns);
+    if (enableMoveInitOperandsToInput) {
+      populateMoveInitOperandsToInputPattern(patterns);
+    }
     (void)applyPatternsGreedily(op, std::move(patterns));
   }
 };

From b4dca9dd33ec57dfd5e6d9ab8de28b61a85f970d Mon Sep 17 00:00:00 2001
From: snigdha dalvi <sdalvi@quicinc.com>
Date: Wed, 10 Sep 2025 13:10:14 -0700
Subject: [PATCH 2/2] Enable MoveInitOperandsToInputPattern for original
 behaviour

---
 mlir/include/mlir/Dialect/Linalg/Passes.td | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlir/include/mlir/Dialect/Linalg/Passes.td b/mlir/include/mlir/Dialect/Linalg/Passes.td
index 7da60bef6d70c..7cd8da6643400 100644
--- a/mlir/include/mlir/Dialect/Linalg/Passes.td
+++ b/mlir/include/mlir/Dialect/Linalg/Passes.td
@@ -147,7 +147,7 @@ def LinalgFoldUnitExtentDimsPass : Pass<"linalg-fold-unit-extent-dims", ""> {
            /*default=*/"false",
            "Generate rank-reducing slices instead of reassociative reshapes">,
     Option<"enableMoveInitOperandsToInput", "enable-move-init-operands-to-input", "bool",
-           /*default=*/"false",
+           /*default=*/"true",
            "Enable MoveInitOperandsToInputPattern transformation">
   ];
   let dependentDialects = [