Disable Attention V operand transposition.

This impacts the ability to horizontally fuse the matmuls that feed into `Q-K-V` transpose. The improvements seen with the change might have been due to reduction in copy overheads, which are no more an issue. Signed-off-by: MaheshRavishankar <[email protected]>
iree-org · Jan 25, 2025 · 417012a · 417012a
1 parent 1648af4
commit 417012a
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 9 deletions.
diff --git a/compiler/src/iree/compiler/GlobalOptimization/Passes.td b/compiler/src/iree/compiler/GlobalOptimization/Passes.td
@@ -115,6 +115,8 @@ def PropagateLinalgTransposePass :
            "Flag used for lit-testing sinking patterns only. Not for general usage">,
     Option<"testBubblingOnly", "test-bubbling-only", "bool", /*default=*/"false",
            "Flag used for lit-testing bubbling patterns only. Not for general usage">,
+    Option<"enableAttentionVTranspose", "enable-attention-v-transpose", "bool",
+            /*default=*/"false", "Enable transposition of attention v operand">,
   ];
 }
 

diff --git a/compiler/src/iree/compiler/GlobalOptimization/PropagateLinalgTranspose.cpp b/compiler/src/iree/compiler/GlobalOptimization/PropagateLinalgTranspose.cpp
@@ -1092,15 +1092,19 @@ void PropagateLinalgTransposePass::runOnOperation() {
     linalg::populateFoldReshapeOpsByExpansionPatterns(bubblingPatterns,
                                                       reshapePropagationFn);
     linalg::FillOp::getCanonicalizationPatterns(bubblingPatterns, context);
-    linalg::ControlFusionFn bubbleTransposeControlFn =
-        [](OpOperand *fusedOperand) {
-          Operation *producer = fusedOperand->get().getDefiningOp();
-          Operation *consumer = fusedOperand->getOwner();
 
-          return IREE::Flow::isNonNullAndOutsideDispatch({producer, consumer});
-        };
-    IREE::LinalgExt::populateBubbleTransposeFromLinalgExtOps(
-        bubblingPatterns, bubbleTransposeControlFn);
+    if (enableAttentionVTranspose) {
+      linalg::ControlFusionFn bubbleTransposeControlFn =
+          [](OpOperand *fusedOperand) {
+            Operation *producer = fusedOperand->get().getDefiningOp();
+            Operation *consumer = fusedOperand->getOwner();
+
+            return IREE::Flow::isNonNullAndOutsideDispatch(
+                {producer, consumer});
+          };
+      IREE::LinalgExt::populateBubbleTransposeFromLinalgExtOps(
+          bubblingPatterns, bubbleTransposeControlFn);
+    }
     bubblingPatterns.insert<FuseTransposeWithProducerLinalgOp>(
         context, enableAggressivePropagation);
     bubblingPatterns.insert<BubbleTransposeThroughCollapseShape>(context);

diff --git a/compiler/src/iree/compiler/GlobalOptimization/test/propagate_linalg_transpose.mlir b/compiler/src/iree/compiler/GlobalOptimization/test/propagate_linalg_transpose.mlir
@@ -1,4 +1,4 @@
-// RUN: iree-opt --pass-pipeline="builtin.module(util.func(iree-global-opt-propagate-linalg-transpose))" --split-input-file %s | FileCheck %s
+// RUN: iree-opt --pass-pipeline="builtin.module(util.func(iree-global-opt-propagate-linalg-transpose{enable-attention-v-transpose=true}))" --split-input-file %s | FileCheck %s
 // RUN: iree-opt --pass-pipeline="builtin.module(util.func(iree-global-opt-propagate-linalg-transpose{enable-aggressive-propagation=true}))" --split-input-file %s | FileCheck %s --check-prefix=APROP
 // RUN: iree-opt --pass-pipeline="builtin.module(util.func(iree-global-opt-propagate-linalg-transpose{test-sinking-only=true}))" --split-input-file %s | FileCheck %s --check-prefix=SINK
 // RUN: iree-opt --pass-pipeline="builtin.module(util.func(iree-global-opt-propagate-linalg-transpose{test-bubbling-only=true}))" --split-input-file %s | FileCheck %s --check-prefix=BUBBLE