[DispatchCreation] Drop unit dims for flow.parameter.named (iree-org#21687)

The input shouldn't have stream.parameter.named attributes at input
level after iree-org#17303 , this patch
makes sure we handle both attributes similarily in DispatchCreation, so
that we see no regressions when frontends switch to this.

Signed-off-by: keshavvinayak01 <[email protected]>

Author: Groverkss

compiler/src/iree/compiler/DispatchCreation/FoldUnitExtentDims.cpp

@@ -11,6 +11,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "iree/compiler/Dialect/Flow/IR/FlowTypes.h"
 #include "iree/compiler/Dialect/Flow/Transforms/RegionOpUtils.h"
 #include "iree/compiler/Dialect/LinalgExt/IR/LinalgExtInterfaces.h"
 #include "iree/compiler/Dialect/LinalgExt/Transforms/Transforms.h"
@@ -252,11 +253,19 @@ foldUnitDimsOnGlobal(IRRewriter &rewriter, IREE::Util::GlobalOpInterface global,
                                                       newGlobalType);
           })
           .Case<IREE::Stream::NamedParameterAttr>(
+              // TODO: Remove this case once frontends have caught up, we should
+              // not have stream.parameter.named at this level.
               [&](IREE::Stream::NamedParameterAttr attr) {
                 return IREE::Stream::NamedParameterAttr::get(
                     rewriter.getContext(), newGlobalType, attr.getScope(),
                     attr.getKey(), attr.getConfig());
               })
+          .Case<IREE::Flow::NamedParameterAttr>(
+              [&](IREE::Flow::NamedParameterAttr attr) {
+                return IREE::Flow::NamedParameterAttr::get(
+                    rewriter.getContext(), newGlobalType, attr.getScope(),
+                    attr.getKey(), attr.getConfig());
+              })
           .Default([&](Attribute) { return nullptr; });
   if (!newInitialValue) {
     return success();

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

@@ -109,6 +109,21 @@ module @fold_stream_parameter {
 
 // -----
 
+module @fold_flow_parameter {
+  util.global private mutable @global = #flow.parameter.named<"module"::"global"> : tensor<1x1x10xf32>
+  util.func public @fold_flow_parameter() -> tensor<1x1x10xf32> {
+    %global = util.global.load @global : tensor<1x1x10xf32>
+    util.return %global : tensor<1x1x10xf32>
+  }
+}
+
+//      CHECK: module @fold_flow_parameter
+//      CHECK:   util.global private mutable @[[GLOBAL:.+]] = #flow.parameter.named<"module"::"global"> : tensor<10xf32>
+//      CHECK:   util.func public @fold_flow_parameter
+//      CHECK:     %[[LOAD:.+]] = util.global.load @[[GLOBAL]] : tensor<10xf32>
+
+// -----
+
 util.func public @scatter(%arg0 : tensor<4xi64>, %arg1 : tensor<4x1xi32>, %arg2 : tensor<4xi64>) -> tensor<4xi64> {
   %0 = iree_linalg_ext.scatter dimension_map = [0] unique_indices(false) ins(%arg0, %arg1: tensor<4xi64>, tensor<4x1xi32>) outs(%arg2 : tensor<4xi64>) {
   ^bb0(%arg3: i64, %arg4: i64):