[DispatchCreation] Use patterns to bubble up expand shape across collapse shapes. (#20648)

Signed-off-by: MaheshRavishankar <mahesh.ravishankar@gmail.com>

Author: MaheshRavishankar

compiler/src/iree/compiler/DispatchCreation/BubbleUpExpandShapes.cpp

@@ -29,6 +29,12 @@
 
 #define DEBUG_TYPE "iree-dispatch-creation-bubble-up-expand-shapes"
 
+static llvm::cl::opt<bool> clPropagateCollapseAcrossExpands(
+    "iree-dispatch-creation-propagate-collapse-across-expands",
+    llvm::cl::desc("Enables change to propagate collapse shapes across expand "
+                   "shapes. This flag is meant as a stop-gap solution before "
+                   "making this default due to codegen issues."),
+    llvm::cl::init(false));
 namespace mlir::iree_compiler::DispatchCreation {
 
 #define GEN_PASS_DEF_BUBBLEUPEXPANDSHAPESPASS
@@ -212,6 +218,10 @@ void BubbleUpExpandShapesPass::runOnOperation() {
   memref::populateResolveRankedShapedTypeResultDimsPatterns(
       bubbleExpandShapePatterns);
 
+  if (clPropagateCollapseAcrossExpands) {
+    tensor::populateBubbleUpExpandShapePatterns(bubbleExpandShapePatterns);
+  }
+
   GreedyRewriteConfig rewriteConfig;
   rewriteConfig.maxIterations = GreedyRewriteConfig::kNoLimit;
   if (failed(applyPatternsGreedily(getOperation(),

compiler/src/iree/compiler/DispatchCreation/test/attention_fuse_by_expansion.mlir

@@ -1,4 +1,4 @@
-// RUN: iree-opt --split-input-file --mlir-print-local-scope --pass-pipeline="builtin.module(util.func(iree-dispatch-creation-bubble-up-expand-shapes, canonicalize, cse, canonicalize))" %s | FileCheck %s
+// RUN: iree-opt --split-input-file --mlir-print-local-scope --pass-pipeline="builtin.module(util.func(iree-dispatch-creation-bubble-up-expand-shapes, canonicalize, cse, canonicalize))" --iree-dispatch-creation-propagate-collapse-across-expands=true %s | FileCheck %s
 
 #map = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2)>
 #map1 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d3, d2)>
@@ -584,8 +584,7 @@ util.func @scatter_collapse_original_partial(%arg0: tensor<?x1x32x8x128xf16>, %a
 //  CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]:
 //  CHECK-SAME:     %[[ARG2:[a-zA-Z0-9]+]]:
 //   CHECK-DAG:   %[[UPDATES:.+]] = tensor.expand_shape %[[ARG0]] {{.*}} tensor<?x1x32x8x128xf16> into tensor<?x1x2x16x4x2x64x2xf16>
-// TODO(IanWood1): fix this so the collapse folds with the expand
-//   CHECK-DAG:   %[[ORIGINAL:.+]] = tensor.expand_shape {{.*}} tensor<?x32x8x128xf16> into tensor<?x2x16x4x2x64x2xf16>
+//   CHECK-DAG:   %[[ORIGINAL:.+]] = tensor.collapse_shape %[[ARG2]] {{.*}} tensor<5x?x2x16x4x2x64x2xf16> into tensor<?x2x16x4x2x64x2xf16>
 //       CHECK:   %[[SCATTER:.+]] = iree_linalg_ext.scatter
 //  CHECK-SAME:       ins(%[[UPDATES]], %[[ARG1]]
 //  CHECK-SAME:       outs(%[[ORIGINAL]]

compiler/src/iree/compiler/Preprocessing/Common/test/attr_based_pipeline.mlir

@@ -1,4 +1,4 @@
-// RUN: iree-opt --iree-preprocessing-attr-based-pipeline --mlir-print-local-scope --split-input-file --verify-diagnostics %s | FileCheck %s
+// RUN: iree-opt --iree-preprocessing-attr-based-pipeline --mlir-print-local-scope --split-input-file --verify-diagnostics --iree-dispatch-creation-propagate-collapse-across-expands=true %s | FileCheck %s
 
 func.func @single_dispatch_dropunitdims(%lhs : tensor<1x26x18x288xbf16>, %rhs :  tensor<288x288x3x3xbf16>, %outs : tensor<1x288x26x18xbf16>,
     %outs2 : tensor<1x288x24x16xf32>) -> tensor<1x288x24x16xf32> attributes {
@@ -12,9 +12,9 @@ func.func @single_dispatch_dropunitdims(%lhs : tensor<1x26x18x288xbf16>, %rhs :
 // CHECK-LABEL: @single_dispatch_dropunitdims
 //  CHECK-SAME: %[[ARG0:[A-Za-z0-9]+]]: tensor<1x26x18x288xbf16>
 //       CHECK:   %[[DISPATCH:.+]] = flow.dispatch.region
-//       CHECK:     %[[COLLAPSE:.+]] = tensor.collapse_shape %[[ARG0]]
-//       CHECK:     %[[EXPAND:.+]] = tensor.expand_shape %[[COLLAPSE]]
-//       CHECK:     %[[CONV:.+]] = linalg.generic {{.*}} ins(%[[EXPAND]]
+//       CHECK:     %[[EXPAND:.+]] = tensor.expand_shape %[[ARG0]]
+//       CHECK:     %[[COLLAPSE:.+]] = tensor.collapse_shape %[[EXPAND]]
+//       CHECK:     %[[CONV:.+]] = linalg.generic {{.*}} ins(%[[COLLAPSE]]
 //       CHECK:     flow.return %[[CONV]]
 //       CHECK:   return %[[DISPATCH]]