[ConstEval] Do not jit parameterized flow.tensor.constants (#21748)
Before, we accepted stream.parameter.named attributes as util.global ops
from the frontend. After we switched to flow.parameter.named ops, some
parameters are now inlined as flow.tensor.constants. Before this patch,
const eval did not treat these ops as parameterized and was trying to
const eval runtime dependent values, which is not possible.
diff --git a/compiler/src/iree/compiler/ConstEval/BUILD.bazel b/compiler/src/iree/compiler/ConstEval/BUILD.bazel
index 1670e29..d370efb 100644
--- a/compiler/src/iree/compiler/ConstEval/BUILD.bazel
+++ b/compiler/src/iree/compiler/ConstEval/BUILD.bazel
@@ -58,6 +58,7 @@
":PassHeaders",
":PassesIncGen",
":Runtime",
+ "//compiler/src/iree/compiler/Dialect/Flow/IR",
"//compiler/src/iree/compiler/Dialect/HAL/Target",
"//compiler/src/iree/compiler/Dialect/Util/Analysis/Constant",
"//compiler/src/iree/compiler/Dialect/Util/IR",
diff --git a/compiler/src/iree/compiler/ConstEval/CMakeLists.txt b/compiler/src/iree/compiler/ConstEval/CMakeLists.txt
index 575085e..54c8792 100644
--- a/compiler/src/iree/compiler/ConstEval/CMakeLists.txt
+++ b/compiler/src/iree/compiler/ConstEval/CMakeLists.txt
@@ -49,6 +49,7 @@
MLIRFunctionInterfaces
MLIRIR
MLIRPass
+ iree::compiler::Dialect::Flow::IR
iree::compiler::Dialect::HAL::Target
iree::compiler::Dialect::Util::Analysis::Constant
iree::compiler::Dialect::Util::IR
diff --git a/compiler/src/iree/compiler/ConstEval/JitGlobals.cpp b/compiler/src/iree/compiler/ConstEval/JitGlobals.cpp
index c49ef30..1467e74 100644
--- a/compiler/src/iree/compiler/ConstEval/JitGlobals.cpp
+++ b/compiler/src/iree/compiler/ConstEval/JitGlobals.cpp
@@ -6,6 +6,7 @@
#include "iree/compiler/ConstEval/Passes.h"
#include "iree/compiler/ConstEval/Runtime.h"
+#include "iree/compiler/Dialect/Flow/IR/FlowOps.h"
#include "iree/compiler/Dialect/HAL/Target/TargetOptions.h"
#include "iree/compiler/Dialect/Util/Analysis/Constant/ConstExpr.h"
#include "iree/compiler/Dialect/Util/Analysis/Constant/OpOracle.h"
@@ -79,6 +80,13 @@
IREEVMPipelineHooks hooks;
};
+static inline bool isAttrParameterized(Attribute attr) {
+ if (!attr)
+ return false;
+ return !isa<IntegerAttr>(attr) && !isa<FloatAttr>(attr) &&
+ !isa<IREE::Util::SerializableAttrInterface>(attr);
+}
+
template <typename AccessorTy>
static inline bool isAccessorParameterized(const SymbolTable &moduleSymbols,
AccessorTy op) {
@@ -86,12 +94,7 @@
moduleSymbols.lookup<IREE::Util::GlobalOpInterface>(op.getGlobalName());
if (!global)
return true;
- auto attr = global.getGlobalInitialValue();
- if (!attr)
- return false;
- return !isa<IntegerAttr>(attr) && !isa<FloatAttr>(attr) &&
- !isa<IREE::Util::SerializableAttrInterface>(
- global.getGlobalInitialValue());
+ return isAttrParameterized(global.getGlobalInitialValue());
}
// Today the only way to interact with a global is with loads, stores, and
@@ -109,6 +112,9 @@
.Case([=](IREE::Util::GlobalStoreOpInterface accessor) {
return isAccessorParameterized(moduleSymbols, accessor);
})
+ .Case([=](IREE::Flow::TensorConstantOp accessor) {
+ return isAttrParameterized(accessor.getValueAttr());
+ })
.Default([=](auto) { return false; });
if (parameterized)
return WalkResult::interrupt();
diff --git a/compiler/src/iree/compiler/ConstEval/test/jit_globals.mlir b/compiler/src/iree/compiler/ConstEval/test/jit_globals.mlir
index 010cbb4..c36cfdd 100644
--- a/compiler/src/iree/compiler/ConstEval/test/jit_globals.mlir
+++ b/compiler/src/iree/compiler/ConstEval/test/jit_globals.mlir
@@ -335,6 +335,20 @@
// -----
+// CHECK-LABEL: @skip_parameterized_tensor_constants
+module @skip_parameterized_tensor_constants {
+ util.global private @hoisted : tensor<f32>
+ // CHECK: util.initializer
+ // expected-warning @+1 {{skipping consteval initializer}}
+ util.initializer {
+ %1 = flow.tensor.constant #flow.parameter.named<"runtime"::"global"> : tensor<f32>
+ util.global.store %1, @hoisted : tensor<f32>
+ util.return
+ }
+}
+
+// -----
+
// TODO(benvanik): rewrite availability to use proper analysis - currently the
// pass uses ConstExprAnalysis which can't actually indicate what we want when
// we want it (here that this iota is available for evaluation at compile-time).
