[DT] Add an option to not hoist infusible encoding ops out of dispatch.

compiler/src/iree/compiler/DispatchCreation/FuseEncodingOpsIntoDispatchRegions.cpp

@@ -42,6 +42,9 @@ struct FuseEncodingOpsIntoDispatchRegionsPass
 
     SmallVector<IREE::Encoding::SetEncodingOp> encodingOps;
     funcOp->walk([&](IREE::Encoding::SetEncodingOp encodingOp) {
+      if (encodingOp->getParentOfType<IREE::Flow::DispatchRegionOp>()) {
+        return;
+      }
       encodingOps.push_back(encodingOp);
     });
 

compiler/src/iree/compiler/DispatchCreation/HoistEncodingOps.cpp

@@ -198,10 +198,17 @@ void HoistEncodingOpsPass::runOnOperation() {
 
   SmallVector<IREE::Encoding::SetEncodingOp> candidates;
   funcOp->walk([&](IREE::Encoding::SetEncodingOp setEncodingOp) {
-    if (setEncodingOp->getParentOfType<IREE::Flow::DispatchRegionOp>()) {
-      candidates.push_back(setEncodingOp);
+    if (!setEncodingOp->getParentOfType<IREE::Flow::DispatchRegionOp>()) {
+      return;
     }
+    if (!hoistInfusibleEncodings &&
+        !setEncodingOp.getSource()
+             .getDefiningOp<IREE::Flow::DispatchRegionOp>()) {
+      return;
+    }
+    candidates.push_back(setEncodingOp);
   });
+
   IRRewriter rewriter(ctx);
   for (auto setEncodingOp : candidates) {
     if (failed(IREE::Flow::hoistOutOfDispatch(rewriter, setEncodingOp))) {

compiler/src/iree/compiler/DispatchCreation/Passes.cpp

@@ -86,6 +86,13 @@ static llvm::cl::opt<bool> clEnableDataTiling(
                    "path, --iree-opt-data-tiling=false must be set as wells"),
     llvm::cl::init(false));
 
+static llvm::cl::opt<bool> clHoistInfusibleEncodings(
+    "iree-dispatch-creation-hoist-infusible-encodings",
+    llvm::cl::desc("Enable the hoisting of infusible encoding ops. The "
+                   "encoding ops may result in their own dispatch. To use this "
+                   "path, --iree-opt-data-tiling=false must be set as wells"),
+    llvm::cl::init(true));
+
 namespace mlir::iree_compiler::DispatchCreation {
 
 //===----------------------------------------------------------------------===//
@@ -253,7 +260,10 @@ static void addDispatchRegionCreationPasses(OpPassManager &passManager) {
         // op, so hoist them out of their current dispatch regions. Also, bubble
         // SetEncodingOps through special operations like bit-extending ops and
         // broadcasting ops.
-        .addPass(DispatchCreation::createHoistEncodingOpsPass)
+        .addPass([&]() {
+          return DispatchCreation::createHoistEncodingOpsPass(
+              HoistEncodingOpsPassOptions{clHoistInfusibleEncodings});
+        })
         // After SetEncodingOps are hoisted, try to fuse them with their
         // producer dispatches to try to hide packing costs.
         .addPass(

compiler/src/iree/compiler/DispatchCreation/Passes.td

@@ -244,6 +244,10 @@ def FuseEncodingOpsIntoDispatchRegionsPass :
 def HoistEncodingOpsPass :
     InterfacePass<"iree-dispatch-creation-hoist-encoding-ops", "mlir::FunctionOpInterface"> {
   let summary = "Hoists tensor encoding ops out of flow dispatch regions.";
+  let options = [
+    Option<"hoistInfusibleEncodings", "hoist-infusible-encodings", "bool", /*default=*/"true",
+           "Unconditionally hoist encoding ops out of dispatch regions.">,
+  ];
   let dependentDialects = [
     "mlir::linalg::LinalgDialect",
     "IREE::Flow::FlowDialect",