Implement ValueBoundsOpInterface on HAL ID/count ops, util.assume.int (#19676)
Add a model of ValueBoundsOpInterface (the one used for reasoning
about affine maps and loops) to the HAL ID ops and util.assume.int to
improve both loop invariant code motion and, in the future,
single-iteration loop removal.
diff --git a/compiler/src/iree/compiler/Codegen/Interfaces/Interfaces.cpp b/compiler/src/iree/compiler/Codegen/Interfaces/Interfaces.cpp
index 8f9f312..e7e0183 100644
--- a/compiler/src/iree/compiler/Codegen/Interfaces/Interfaces.cpp
+++ b/compiler/src/iree/compiler/Codegen/Interfaces/Interfaces.cpp
@@ -64,6 +64,7 @@
affine::registerValueBoundsOpInterfaceExternalModels(registry);
arith::registerValueBoundsOpInterfaceExternalModels(registry);
bufferization::registerTransformDialectExtension(registry);
+ gpu::registerValueBoundsOpInterfaceExternalModels(registry);
gpu::registerTransformDialectExtension(registry);
gpu::registerValueBoundsOpInterfaceExternalModels(registry);
linalg::registerTransformDialectExtension(registry);
diff --git a/compiler/src/iree/compiler/ExternalInterfaces/BUILD.bazel b/compiler/src/iree/compiler/ExternalInterfaces/BUILD.bazel
index ae0370c..40a9fe4 100644
--- a/compiler/src/iree/compiler/ExternalInterfaces/BUILD.bazel
+++ b/compiler/src/iree/compiler/ExternalInterfaces/BUILD.bazel
@@ -16,6 +16,7 @@
name = "ExternalModels",
srcs = [
"FlowExternalModels.cpp",
+ "HALExternalModels.cpp",
"Interfaces.cpp",
"LinalgExtExternalModels.cpp",
"StreamExternalModels.cpp",
@@ -23,6 +24,7 @@
],
hdrs = [
"FlowExternalModels.h",
+ "HALExternalModels.h",
"Interfaces.h",
"LinalgExtExternalModels.h",
"StreamExternalModels.h",
diff --git a/compiler/src/iree/compiler/ExternalInterfaces/CMakeLists.txt b/compiler/src/iree/compiler/ExternalInterfaces/CMakeLists.txt
index f7ef47e..18e3522 100644
--- a/compiler/src/iree/compiler/ExternalInterfaces/CMakeLists.txt
+++ b/compiler/src/iree/compiler/ExternalInterfaces/CMakeLists.txt
@@ -15,12 +15,14 @@
ExternalModels
HDRS
"FlowExternalModels.h"
+ "HALExternalModels.h"
"Interfaces.h"
"LinalgExtExternalModels.h"
"StreamExternalModels.h"
"UtilExternalModels.h"
SRCS
"FlowExternalModels.cpp"
+ "HALExternalModels.cpp"
"Interfaces.cpp"
"LinalgExtExternalModels.cpp"
"StreamExternalModels.cpp"
diff --git a/compiler/src/iree/compiler/ExternalInterfaces/HALExternalModels.cpp b/compiler/src/iree/compiler/ExternalInterfaces/HALExternalModels.cpp
new file mode 100644
index 0000000..375368b
--- /dev/null
+++ b/compiler/src/iree/compiler/ExternalInterfaces/HALExternalModels.cpp
@@ -0,0 +1,66 @@
+// Copyright 2025 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "iree/compiler/ExternalInterfaces/HALExternalModels.h"
+
+#include "iree/compiler/Dialect/HAL/IR/HALDialect.h"
+#include "iree/compiler/Dialect/HAL/IR/HALOps.h"
+#include "mlir/Interfaces/ValueBoundsOpInterface.h"
+
+namespace mlir::iree_compiler {
+
+namespace {
+
+//===----------------------------------------------------------------------===//
+// ValueBoundsOpInterface
+//===----------------------------------------------------------------------===//
+
+template <typename IDOp>
+struct IDOpValueBoundsInterface : public ValueBoundsOpInterface::ExternalModel<
+ IDOpValueBoundsInterface<IDOp>, IDOp> {
+ void populateBoundsForIndexValue(Operation *op, Value value,
+ ValueBoundsConstraintSet &cstr) const {
+ auto boundOp = cast<IDOp>(op);
+ assert(value == boundOp.getResult() && "value must be op result");
+ cstr.bound(value) >= 0;
+ if (boundOp.getUpperBound()) {
+ cstr.bound(value) < boundOp.getUpperBound()->getSExtValue();
+ }
+ }
+};
+
+template <typename CountOp>
+struct CountOpValueBoundsInterface
+ : public ValueBoundsOpInterface::ExternalModel<
+ CountOpValueBoundsInterface<CountOp>, CountOp> {
+ void populateBoundsForIndexValue(Operation *op, Value value,
+ ValueBoundsConstraintSet &cstr) const {
+ auto boundOp = cast<CountOp>(op);
+ assert(value == boundOp.getResult() && "value must be op result");
+ cstr.bound(value) >= 1;
+ if (boundOp.getUpperBound()) {
+ cstr.bound(value) <= boundOp.getUpperBound()->getSExtValue();
+ }
+ }
+};
+
+} // namespace
+
+void registerHALExternalModels(DialectRegistry ®istry) {
+ registry.addExtension(+[](MLIRContext *context,
+ IREE::HAL::HALDialect *dialect) {
+ IREE::HAL::InterfaceWorkgroupIDOp::attachInterface<
+ IDOpValueBoundsInterface<IREE::HAL::InterfaceWorkgroupIDOp>>(*context);
+
+ IREE::HAL::InterfaceWorkgroupSizeOp::attachInterface<
+ CountOpValueBoundsInterface<IREE::HAL::InterfaceWorkgroupSizeOp>>(
+ *context);
+ IREE::HAL::InterfaceWorkgroupCountOp::attachInterface<
+ CountOpValueBoundsInterface<IREE::HAL::InterfaceWorkgroupCountOp>>(
+ *context);
+ });
+}
+} // namespace mlir::iree_compiler
diff --git a/compiler/src/iree/compiler/ExternalInterfaces/HALExternalModels.h b/compiler/src/iree/compiler/ExternalInterfaces/HALExternalModels.h
new file mode 100644
index 0000000..9bc3c34
--- /dev/null
+++ b/compiler/src/iree/compiler/ExternalInterfaces/HALExternalModels.h
@@ -0,0 +1,20 @@
+// Copyright 2025 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#ifndef IREE_COMPILER_EXTERNALINTERFACES_HALEXTERNALMODELS_H_
+#define IREE_COMPILER_EXTERNALINTERFACES_HALEXTERNALMODELS_H_
+
+namespace mlir {
+class DialectRegistry;
+} // namespace mlir
+
+namespace mlir::iree_compiler {
+
+void registerHALExternalModels(DialectRegistry ®istry);
+
+} // namespace mlir::iree_compiler
+
+#endif // IREE_COMPILER_EXTERNALINTERFACES_HALEXTERNALMODELS_H_
diff --git a/compiler/src/iree/compiler/ExternalInterfaces/Interfaces.cpp b/compiler/src/iree/compiler/ExternalInterfaces/Interfaces.cpp
index 2a6ffa3..fc7d0fe 100644
--- a/compiler/src/iree/compiler/ExternalInterfaces/Interfaces.cpp
+++ b/compiler/src/iree/compiler/ExternalInterfaces/Interfaces.cpp
@@ -7,6 +7,7 @@
#include "iree/compiler/ExternalInterfaces/Interfaces.h"
#include "iree/compiler/ExternalInterfaces/FlowExternalModels.h"
+#include "iree/compiler/ExternalInterfaces/HALExternalModels.h"
#include "iree/compiler/ExternalInterfaces/LinalgExtExternalModels.h"
#include "iree/compiler/ExternalInterfaces/StreamExternalModels.h"
#include "iree/compiler/ExternalInterfaces/UtilExternalModels.h"
@@ -15,6 +16,7 @@
void registerExternalInterfaces(DialectRegistry ®istry) {
registerFlowExternalModels(registry);
+ registerHALExternalModels(registry);
registerLinalgExtExternalModels(registry);
registerStreamExternalModels(registry);
registerUtilExternalModels(registry);
diff --git a/compiler/src/iree/compiler/ExternalInterfaces/UtilExternalModels.cpp b/compiler/src/iree/compiler/ExternalInterfaces/UtilExternalModels.cpp
index db5fa40..1990dac 100644
--- a/compiler/src/iree/compiler/ExternalInterfaces/UtilExternalModels.cpp
+++ b/compiler/src/iree/compiler/ExternalInterfaces/UtilExternalModels.cpp
@@ -12,12 +12,15 @@
#include "iree/compiler/Dialect/Encoding/IR/EncodingOps.h"
#include "iree/compiler/Dialect/LinalgExt/IR/LinalgExtDialect.h"
#include "iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.h"
+#include "iree/compiler/Dialect/Util/IR/UtilDialect.h"
+#include "iree/compiler/Dialect/Util/IR/UtilOps.h"
#include "iree/compiler/Dialect/Util/IR/UtilTypes.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/MLProgram/IR/MLProgram.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/IR/Matchers.h"
+#include "mlir/Interfaces/ValueBoundsOpInterface.h"
namespace mlir::iree_compiler {
@@ -105,6 +108,31 @@
};
//===----------------------------------------------------------------------===//
+// ValueBoundsOpInterface
+//===----------------------------------------------------------------------===//
+
+/// For some reason, this interface has to be done as an external model.
+struct UtilAssumeIntValueBoundsOpInterface
+ : public ValueBoundsOpInterface::ExternalModel<
+ UtilAssumeIntValueBoundsOpInterface, IREE::Util::AssumeIntOp> {
+ void populateBoundsForIndexValue(Operation *op, Value value,
+ ValueBoundsConstraintSet &cstr) const {
+ auto assumeOp = cast<IREE::Util::AssumeIntOp>(op);
+ auto result = cast<OpResult>(value);
+ assert(result.getOwner() == op && "value is a result of this index op");
+ auto [min, max] =
+ assumeOp.getUnionedUnsignedRange(result.getResultNumber());
+
+ if (min) {
+ cstr.bound(result) >= *min;
+ }
+ if (max) {
+ cstr.bound(result) <= *max;
+ }
+ }
+};
+
+//===----------------------------------------------------------------------===//
// GlobalOpInterface
//===----------------------------------------------------------------------===//
@@ -489,6 +517,11 @@
tensor::PadOp, tensor::PackOp,
tensor::UnPackOp>::registerOpInterface(context);
});
+ registry.addExtension(
+ +[](MLIRContext *context, IREE::Util::UtilDialect *dialect) {
+ IREE::Util::AssumeIntOp::attachInterface<
+ UtilAssumeIntValueBoundsOpInterface>(*context);
+ });
}
} // namespace mlir::iree_compiler
