Implement ValueBoundsOpInterface on HAL ID/count ops, util.assume.int (iree-org#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.

Author: krzysz00

compiler/src/iree/compiler/Codegen/Interfaces/Interfaces.cpp

@@ -64,6 +64,7 @@ void registerCodegenInterfaces(DialectRegistry &registry) {
   affine::registerValueBoundsOpInterfaceExternalModels(registry);
   arith::registerValueBoundsOpInterfaceExternalModels(registry);
   bufferization::registerTransformDialectExtension(registry);
+  gpu::registerValueBoundsOpInterfaceExternalModels(registry);
   gpu::registerTransformDialectExtension(registry);
   gpu::registerValueBoundsOpInterfaceExternalModels(registry);
   linalg::registerTransformDialectExtension(registry);

compiler/src/iree/compiler/ExternalInterfaces/BUILD.bazel

@@ -16,13 +16,15 @@ iree_compiler_cc_library(
     name = "ExternalModels",
     srcs = [
         "FlowExternalModels.cpp",
+        "HALExternalModels.cpp",
         "Interfaces.cpp",
         "LinalgExtExternalModels.cpp",
         "StreamExternalModels.cpp",
         "UtilExternalModels.cpp",
     ],
     hdrs = [
         "FlowExternalModels.h",
+        "HALExternalModels.h",
         "Interfaces.h",
         "LinalgExtExternalModels.h",
         "StreamExternalModels.h",

compiler/src/iree/compiler/ExternalInterfaces/CMakeLists.txt

@@ -15,12 +15,14 @@ iree_cc_library(
     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"

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 &registry) {
+  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

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 &registry);
+
+} // namespace mlir::iree_compiler
+
+#endif // IREE_COMPILER_EXTERNALINTERFACES_HALEXTERNALMODELS_H_

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 @@ namespace mlir::iree_compiler {
 
 void registerExternalInterfaces(DialectRegistry &registry) {
   registerFlowExternalModels(registry);
+  registerHALExternalModels(registry);
   registerLinalgExtExternalModels(registry);
   registerStreamExternalModels(registry);
   registerUtilExternalModels(registry);

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 {
 
@@ -104,6 +107,31 @@ struct ArithDivUIInferIntDivisibilityOpInterface
   }
 };
 
+//===----------------------------------------------------------------------===//
+// 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 @@ void registerUtilExternalModels(DialectRegistry &registry) {
             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