[DT] Add a flag to not hoist encodings when the source is ConstExpr. (#20526)

The revision adds an option to not hoist encodings when the source is a
global.load op or has ConstantLike trait.

It is a workaround for huge memory allocation because the ops are still
in dispatch region. I.e., it stops hoisting globals to initializer.

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

Author: hanhanW

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

@@ -10,6 +10,7 @@
 #include "iree/compiler/Dialect/Flow/IR/FlowOps.h"
 #include "iree/compiler/Dialect/Flow/Transforms/RegionOpUtils.h"
 #include "iree/compiler/Dialect/LinalgExt/Utils/Utils.h"
+#include "iree/compiler/Dialect/Util/IR/UtilOps.h"
 #include "iree/compiler/DispatchCreation/Passes.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/Support/Debug.h"
@@ -188,9 +189,16 @@ 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;
     }
+    Operation *src = setEncodingOp.getSource().getDefiningOp();
+    if (!hoistEncodingsForConstExpr && src &&
+        (isa<IREE::Util::GlobalLoadOp>(src) ||
+         src->hasTrait<OpTrait::ConstantLike>())) {
+      return;
+    }
+    candidates.push_back(setEncodingOp);
   });
   IRRewriter rewriter(ctx);
   for (auto setEncodingOp : candidates) {

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

@@ -80,6 +80,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> clHoistEncodingsForConstExpr(
+    "iree-dispatch-creation-hoist-encodings-for-constexpr",
+    llvm::cl::desc("Enable the hoisting of encoding ops when the source is "
+                   "from globals. To use this path, "
+                   "--iree-opt-data-tiling=false must be set as wells"),
+    llvm::cl::init(true));
+
 namespace mlir::iree_compiler::DispatchCreation {
 
 //===----------------------------------------------------------------------===//
@@ -252,9 +259,10 @@ 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)
-        // After SetEncodingOps are hoisted, try to fuse them with their
-        // producer dispatches to try to hide packing costs.
+        .addPass([&]() {
+          return DispatchCreation::createHoistEncodingOpsPass(
+              HoistEncodingOpsPassOptions{clHoistEncodingsForConstExpr});
+        })
         .addPass(
             DispatchCreation::createFuseEncodingOpsIntoDispatchRegionsPass);
   }

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

@@ -286,6 +286,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<"hoistEncodingsForConstExpr", "hoist-encodings-for-constexpr", "bool", /*default=*/"true",
+           "Enable the hoisting when the source is a ConstExpr.">,
+  ];
   let dependentDialects = [
     "mlir::linalg::LinalgDialect",
     "IREE::Flow::FlowDialect",

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

@@ -1,4 +1,5 @@
 // RUN: iree-opt --pass-pipeline="builtin.module(util.func(iree-dispatch-creation-hoist-encoding-ops))" --split-input-file %s | FileCheck %s
+// RUN: iree-opt --pass-pipeline="builtin.module(util.func(iree-dispatch-creation-hoist-encoding-ops{hoist-encodings-for-constexpr=false}))" --split-input-file %s | FileCheck %s --check-prefix=NO-CONST
 
 #map1 = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
 #map2 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
@@ -247,3 +248,8 @@ util.func public @hoist_encoding_only() -> tensor<640x320xf32> {
 // CHECK:         %[[CST:.+]] = arith.constant
 // CHECK:         %[[SET_ENCODING:.+]] = iree_encoding.set_encoding %[[CST]]
 // CHECK:         flow.dispatch.region
+
+// NO-CONST-LABEL: util.func public @hoist_encoding_only(
+// NO-CONST:         %[[CST:.+]] = arith.constant
+// NO-CONST:         flow.dispatch.region
+// NO-CONST:           %[[SET_ENCODING:.+]] = iree_encoding.set_encoding %[[CST]]