[Mlir-commits] [mlir] Flag to guard MoveInitOperandsToInput in LinalgFoldUnitExtendDimsPass (PR #157941)

[llvmlistbot at llvm.org]

https://github.com/sdalvi-quic created https://github.com/llvm/llvm-project/pull/157941

The pattern **MoveInitOperandsToInput** in the LinalgFoldUnitExtendDimsPass tries to move the init operands to input, to ensure that the element wise ops are canonicalized. I need to run **split-reduction** optimization before **LinalgFoldUnitExtendDimsPass**. The split reduction optimization is trying to optimize the Linalg op by reading from the destination operand but, the pattern MoveInitOperandsToInput is reverting it and trying to canonicalize element wise op. This is leading to regression in performance. Proposing to guard the pattern MoveInitOperandsToInput to have more control over enabling/disabling it. 

IR after Split-Reduction and before LinalgFoldUnitExtendDimsPass
<img width="1243" height="290" alt="image" src="https://github.com/user-attachments/assets/4c9a08e3-3d89-4fa8-a484-18eff21eeadb" />

IR after LinalgFoldUnitExtendDimsPass
<img width="1246" height="279" alt="image" src="https://github.com/user-attachments/assets/91b4ddb9-a73c-4f0c-965e-0616066a4b71" />


>From e7061c59c1cb5e37d1bf9dfb5c7b8529b40cd9f6 Mon Sep 17 00:00:00 2001
From: snigdha dalvi <sdalvi at quicinc.com>
Date: Wed, 10 Sep 2025 12:28:33 -0700
Subject: [PATCH] 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));
   }
 };