[Flow] Loosen restrictions on body ops of dequant-like ops (#16449)

The restrictions for selecting dequant-like ops were too conservative.
It should not matter what the body of the operation is as long as the
primary input bitwidth is decreased.

compiler/src/iree/compiler/Dialect/Flow/Transforms/FormDispatchRegions.cpp

diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/FormDispatchRegions.cpp b/compiler/src/iree/compiler/Dialect/Flow/Transforms/FormDispatchRegions.cpp
index aa88c28..652c7eb 100644
--- a/compiler/src/iree/compiler/Dialect/Flow/Transforms/FormDispatchRegions.cpp
+++ b/compiler/src/iree/compiler/Dialect/Flow/Transforms/FormDispatchRegions.cpp
@@ -218,6 +218,10 @@
   if (op->getParentOfType<IREE::Flow::DispatchWorkgroupsOp>()) {
     return false;
   }
+  // Dequantization-like ops get cloned into dispatches later.
+  if (isDequantizationLikeOp(op)) {
+    return false;
+  }
   // Any Linalg named op or generic op with reduction iterator types is a root
   // op.
   if (auto linalgOp = dyn_cast<linalg::LinalgOp>(op)) {

compiler/src/iree/compiler/Dialect/Flow/Transforms/RegionOpUtils.cpp

diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/RegionOpUtils.cpp b/compiler/src/iree/compiler/Dialect/Flow/Transforms/RegionOpUtils.cpp
index 6129744..a130997 100644
--- a/compiler/src/iree/compiler/Dialect/Flow/Transforms/RegionOpUtils.cpp
+++ b/compiler/src/iree/compiler/Dialect/Flow/Transforms/RegionOpUtils.cpp
@@ -556,15 +556,6 @@
       outputElementType.getIntOrFloatBitWidth()) {
     return false;
   }
-
-  for (auto &bodyOp : genericOp.getBody()->getOperations()) {
-    if (!isa<arith::ExtUIOp, arith::ExtSIOp, arith::MulFOp, arith::MulIOp,
-             arith::AddFOp, arith::AddIOp, arith::SubFOp, arith::SubIOp,
-             arith::SIToFPOp, arith::UIToFPOp, linalg::YieldOp>(bodyOp)) {
-      return false;
-    }
-  }
-
   return true;
 }
 

compiler/src/iree/compiler/Dialect/Flow/Transforms/test/clone_producers_into_dispatch_regions.mlir

diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/clone_producers_into_dispatch_regions.mlir b/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/clone_producers_into_dispatch_regions.mlir
index 136af8b..20b3dc2 100644
--- a/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/clone_producers_into_dispatch_regions.mlir
+++ b/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/clone_producers_into_dispatch_regions.mlir
@@ -237,3 +237,34 @@
 //  CHECK-SAME:       outs(%[[FILL]] :
 //       CHECK:   flow.return %[[MMT4D]] :
 //       CHECK:   util.return %[[DISP]]
+
+// -----
+
+#map = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+#map1 = affine_map<(d0, d1, d2) -> (d0)>
+util.func public @dequant_like_extf_reduction(%arg0: tensor<11008x32x128xf16>) -> tensor<11008xf32> {
+  %0 = tensor.empty() : tensor<11008x32x128xf32>
+  %1 = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel", "parallel", "parallel"]} ins(%arg0 : tensor<11008x32x128xf16>) outs(%0 : tensor<11008x32x128xf32>) {
+  ^bb0(%in: f16, %out: f32):
+    %4 = arith.extf %in : f16 to f32
+    linalg.yield %4 : f32
+  } -> tensor<11008x32x128xf32>
+  %2 = tensor.empty() : tensor<11008xf32>
+  %3 = flow.dispatch.region -> (tensor<11008xf32>) {
+    %4 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "reduction", "reduction"]} ins(%1 : tensor<11008x32x128xf32>) outs(%2 : tensor<11008xf32>) {
+    ^bb0(%in: f32, %out: f32):
+      %5 = arith.addf %in, %out : f32
+      linalg.yield %5 : f32
+    } -> tensor<11008xf32>
+    flow.return %4 : tensor<11008xf32>
+  }
+  util.return %3 : tensor<11008xf32>
+}
+//       CHECK: util.func public @dequant_like_extf_reduction
+//       CHECK:   %[[DISP:.+]] = flow.dispatch.region -> (tensor<11008xf32>)
+//       CHECK:   linalg.generic
+//  CHECK-SAME:       iterator_types = ["parallel", "parallel", "parallel"]
+//       CHECK:   %[[GEN:.+]] = linalg.generic
+//  CHECK-SAME:       iterator_types = ["parallel", "reduction", "reduction"]
+//       CHECK:   flow.return %[[GEN]] :
+//       CHECK:   util.return %[[DISP]]