From f66b49af61b75fce32616f465c44ac52b59fdf41 Mon Sep 17 00:00:00 2001
From: Hanhan Wang <hanchung@google.com>
Date: Wed, 23 Aug 2023 15:47:43 -0700
Subject: [PATCH] Disable hoisting target-dependent ops to globals

---
 .../compiler/Dialect/Util/Analysis/Constant/BUILD.bazel   | 2 ++
 .../Dialect/Util/Analysis/Constant/CMakeLists.txt         | 2 ++
 .../compiler/Dialect/Util/Analysis/Constant/OpOracle.cpp  | 8 ++++++++
 3 files changed, 12 insertions(+)

diff --git a/compiler/src/iree/compiler/Dialect/Util/Analysis/Constant/BUILD.bazel b/compiler/src/iree/compiler/Dialect/Util/Analysis/Constant/BUILD.bazel
index 47bcaa17ab19f..58c86e1cf4305 100644
--- a/compiler/src/iree/compiler/Dialect/Util/Analysis/Constant/BUILD.bazel
+++ b/compiler/src/iree/compiler/Dialect/Util/Analysis/Constant/BUILD.bazel
@@ -23,8 +23,10 @@ iree_compiler_cc_library(
         "OpOracle.h",
     ],
     deps = [
+        "//compiler/src/iree/compiler/Codegen/Dialect:IREECodegenDialect",
         "//compiler/src/iree/compiler/Dialect/Util/Analysis",
         "//compiler/src/iree/compiler/Dialect/Util/IR",
+        "//llvm-external-projects/iree-dialects:IREELinalgExtDialect",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:ArithDialect",
         "@llvm-project//mlir:FuncDialect",
diff --git a/compiler/src/iree/compiler/Dialect/Util/Analysis/Constant/CMakeLists.txt b/compiler/src/iree/compiler/Dialect/Util/Analysis/Constant/CMakeLists.txt
index 9fe654f320ef7..48bfef0585ff3 100644
--- a/compiler/src/iree/compiler/Dialect/Util/Analysis/Constant/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Dialect/Util/Analysis/Constant/CMakeLists.txt
@@ -20,6 +20,7 @@ iree_cc_library(
     "ConstExpr.cpp"
     "OpOracle.cpp"
   DEPS
+    IREELinalgExtDialect
     LLVMSupport
     MLIRArithDialect
     MLIRFuncDialect
@@ -27,6 +28,7 @@ iree_cc_library(
     MLIRLinalgDialect
     MLIRSupport
     MLIRTensorDialect
+    iree::compiler::Codegen::Dialect::IREECodegenDialect
     iree::compiler::Dialect::Util::Analysis
     iree::compiler::Dialect::Util::IR
   PUBLIC
diff --git a/compiler/src/iree/compiler/Dialect/Util/Analysis/Constant/OpOracle.cpp b/compiler/src/iree/compiler/Dialect/Util/Analysis/Constant/OpOracle.cpp
index 5b812b4c614af..6351073257a54 100644
--- a/compiler/src/iree/compiler/Dialect/Util/Analysis/Constant/OpOracle.cpp
+++ b/compiler/src/iree/compiler/Dialect/Util/Analysis/Constant/OpOracle.cpp
@@ -6,6 +6,8 @@
 
 #include "iree/compiler/Dialect/Util/Analysis/Constant/OpOracle.h"
 
+#include "iree-dialects/Dialect/LinalgExt/IR/LinalgExtOps.h"
+#include "iree/compiler/Codegen/Dialect/IREECodegenOps.h"
 #include "iree/compiler/Dialect/Util/IR/UtilDialect.h"
 #include "iree/compiler/Dialect/Util/IR/UtilOps.h"
 #include "llvm/ADT/SmallPtrSet.h"
@@ -100,6 +102,12 @@ ConstExprOpInfo ConstExprOpInfo::getForOp(Operation *op) {
     return {};
   }
 
+  // Target-dependent ops are not const-expr.
+  if (isa<IREE::LinalgExt::UpperBoundTileSizeOp, IREE::LinalgExt::SetEncodingOp,
+          IREE::Codegen::QueryTileSizesOp>(op)) {
+    return {};
+  }
+
   // By default, ops without results are not const-expr.
   if (op->getNumResults() == 0) {
     return {};