[Preprocessing] Add op for matching dimension bounds (#20502)

Currently there is only support for matching static shapes and dimension
alignment. This allows matching ops based on some minimum/maximum size
for problem shape sensitive strategies.

Author: qedawkins

compiler/src/iree/compiler/Preprocessing/Common/test/preprocessing_match_ops.mlir

@@ -139,6 +139,83 @@ module attributes {transform.with_named_sequence} {
 
 // -----
 
+func.func private @external(%arg0: tensor<?xf32>)
+func.func private @external_lb(%arg0: tensor<100xf32>)
+func.func private @external_ub(%arg0: tensor<3xf32>)
+func.func private @external_2d(%arg0: tensor<?x20xf32>)
+
+// CHECK-LABEL: func.func @call_external
+func.func @call_external(%arg0: index,
+                         %input_2d: tensor<?x20xf32>,
+                         %input_lb: tensor<100xf32>,
+                         %input_ub: tensor<3xf32>) {
+%0 = util.assume.int %arg0<umin = 12, umax = 16, udiv = 1> : index
+%input = tensor.empty(%0) : tensor<?xf32>
+//       CHECK: call @external
+//  CHECK-SAME:   match_status = "both_matched"
+  func.call @external(%input) {match_status = "unmatched"} : (tensor<?xf32>) -> ()
+//       CHECK: call @external_2d
+//  CHECK-SAME:   match_status = "dim1_matched"
+  func.call @external_2d(%input_2d) {match_status = "unmatched"} : (tensor<?x20xf32>) -> ()
+//       CHECK: call @external_lb
+//  CHECK-SAME:   match_status = "lb_matched"
+  func.call @external_lb(%input_lb) {match_status = "unmatched"} : (tensor<100xf32>) -> ()
+//       CHECK: call @external_ub
+//  CHECK-SAME:   match_status = "ub_matched"
+  func.call @external_ub(%input_ub) {match_status = "unmatched"} : (tensor<3xf32>) -> ()
+  return
+}
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @dim1_match(%call: !transform.any_op {transform.readonly}) -> (!transform.any_op, !transform.any_param) {
+    transform.match.operation_name %call ["func.call"] : !transform.any_op
+    %in0 = transform.get_operand %call[0] : (!transform.any_op) -> !transform.any_value
+    transform.iree.match.dim_bounds %in0[1], umin = 20, umax = 20 : !transform.any_value
+    %0 = transform.param.constant "dim1_matched" -> !transform.any_param
+    transform.yield %call, %0 : !transform.any_op, !transform.any_param
+  }
+  transform.named_sequence @both_match(%call: !transform.any_op {transform.readonly}) -> (!transform.any_op, !transform.any_param) {
+    transform.match.operation_name %call ["func.call"] : !transform.any_op
+    %in0 = transform.get_operand %call[0] : (!transform.any_op) -> !transform.any_value
+    transform.iree.match.dim_bounds %in0[0], umin = 5, umax = 20 : !transform.any_value
+    %0 = transform.param.constant "both_matched" -> !transform.any_param
+    transform.yield %call, %0 : !transform.any_op, !transform.any_param
+  }
+  transform.named_sequence @lb_match(%call: !transform.any_op {transform.readonly}) -> (!transform.any_op, !transform.any_param) {
+    transform.match.operation_name %call ["func.call"] : !transform.any_op
+    %in0 = transform.get_operand %call[0] : (!transform.any_op) -> !transform.any_value
+    transform.iree.match.dim_bounds %in0[0], umin = 75, none : !transform.any_value
+    %0 = transform.param.constant "lb_matched" -> !transform.any_param
+    transform.yield %call, %0 : !transform.any_op, !transform.any_param
+  }
+  transform.named_sequence @ub_match(%call: !transform.any_op {transform.readonly}) -> (!transform.any_op, !transform.any_param) {
+    transform.match.operation_name %call ["func.call"] : !transform.any_op
+    %in0 = transform.get_operand %call[0] : (!transform.any_op) -> !transform.any_value
+    transform.iree.match.dim_bounds %in0[0], none, umax = 4 : !transform.any_value
+    %0 = transform.param.constant "ub_matched" -> !transform.any_param
+    transform.yield %call, %0 : !transform.any_op, !transform.any_param
+  }
+
+  transform.named_sequence @annotate(%call: !transform.any_op {transform.readonly},
+                                     %note: !transform.any_param {transform.readonly}) {
+    transform.annotate %call "match_status" = %note : !transform.any_op, !transform.any_param
+    transform.yield
+  }
+
+  transform.named_sequence @__transform_main(%module: !transform.any_op) {
+    %func = transform.structured.match ops{["func.func"]} in %module : (!transform.any_op) -> !transform.any_op
+    transform.foreach_match in %module
+        @dim1_match -> @annotate,
+        @both_match -> @annotate,
+        @lb_match -> @annotate,
+        @ub_match -> @annotate
+      : (!transform.any_op) -> (!transform.any_op)
+    transform.yield
+  }
+}
+
+// -----
+
 module attributes {transform.with_named_sequence} {
 
   // CHECK: func.func @matmul_repeated_operand

compiler/src/iree/compiler/Preprocessing/TransformExtensions/BUILD.bazel

@@ -66,5 +66,6 @@ iree_compiler_cc_library(
         "@llvm-project//mlir:TransformDialect",
         "@llvm-project//mlir:TransformDialectInterfaces",
         "@llvm-project//mlir:TransformUtils",
+        "@llvm-project//mlir:ValueBoundsOpInterface",
     ],
 )

compiler/src/iree/compiler/Preprocessing/TransformExtensions/CMakeLists.txt

@@ -40,6 +40,7 @@ iree_cc_library(
     MLIRTransformDialect
     MLIRTransformDialectInterfaces
     MLIRTransformUtils
+    MLIRValueBoundsOpInterface
     iree::compiler::Utils
   PUBLIC
 )

compiler/src/iree/compiler/Preprocessing/TransformExtensions/PreprocessingExtensions.cpp

@@ -14,6 +14,7 @@
 #include "mlir/IR/IRMapping.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/Value.h"
+#include "mlir/Interfaces/ValueBoundsOpInterface.h"
 
 namespace mlir::iree_compiler {
 
@@ -271,6 +272,53 @@ IREE::transform_dialect::MatchCastCompatibleTypesOp::matchValue(
   return DiagnosedSilenceableFailure::success();
 }
 
+//===----------------------------------------------------------------------===//
+// MatchDimBoundsOp
+//===----------------------------------------------------------------------===//
+
+DiagnosedSilenceableFailure
+IREE::transform_dialect::MatchDimBoundsOp::matchValue(
+    Value current, transform::TransformResults &results,
+    transform::TransformState &state) {
+  auto shapedType = dyn_cast<ShapedType>(current.getType());
+  if (!shapedType) {
+    return emitSilenceableError()
+           << "type " << current.getType() << " is not a shaped type";
+  }
+  int64_t dim = getDim();
+  if (dim >= shapedType.getRank()) {
+    return emitSilenceableError()
+           << "dim " << dim << " out of range for shaped type " << shapedType;
+  }
+  if (std::optional<int64_t> lb = getLowerBound()) {
+    auto constantLb = ValueBoundsConstraintSet::computeConstantBound(
+        presburger::BoundType::LB, {current, /*dim=*/dim},
+        /*stopCondition=*/nullptr, /*closedLB=*/true);
+    if (failed(constantLb)) {
+      return emitSilenceableError()
+             << "failed to compute constant lower bound for dim " << dim;
+    }
+    if (lb.value() > constantLb.value()) {
+      return emitSilenceableError()
+             << "dim " << dim << " is not >= " << lb.value();
+    }
+  }
+  if (std::optional<int64_t> ub = getUpperBound()) {
+    auto constantUb = ValueBoundsConstraintSet::computeConstantBound(
+        presburger::BoundType::UB, {current, /*dim=*/dim},
+        /*stopCondition=*/nullptr, /*closedUB=*/true);
+    if (failed(constantUb)) {
+      return emitSilenceableError()
+             << "failed to compute constant upper bound for dim " << dim;
+    }
+    if (ub.value() < constantUb.value()) {
+      return emitSilenceableError()
+             << "dim " << dim << " is not <= " << ub.value();
+    }
+  }
+  return DiagnosedSilenceableFailure::success();
+}
+
 //===----------------------------------------------------------------------===//
 // MatchDimIsMultipleOfOp
 //===----------------------------------------------------------------------===//
@@ -285,7 +333,7 @@ IREE::transform_dialect::MatchDimIsMultipleOfOp::matchValue(
            << "type " << current.getType() << " is not a shaped type";
   }
   int64_t dim = getDim();
-  if (dim > shapedType.getRank()) {
+  if (dim >= shapedType.getRank()) {
     return emitSilenceableError()
            << "dim " << dim << " out of range for shaped type " << shapedType;
   }

compiler/src/iree/compiler/Preprocessing/TransformExtensions/PreprocessingExtensionsOps.td

@@ -86,13 +86,45 @@ def MatchCastCompatibleTypesOp : Op<Transform_Dialect, "iree.match.cast_compatib
   let cppNamespace = "mlir::iree_compiler::IREE::transform_dialect";
 }
 
+def MatchDimBoundsOp : Op<Transform_Dialect, "iree.match.dim_bounds",
+    [IsolatedFromAbove,
+     MatchOpInterface,
+     SingleValueMatcher,
+     MemoryEffectsOpInterface]> {
+  let summary =
+      "Checks whether the size of a dim is within given bounds";
+  let description = [{
+    Checks whether a dim is within a specified lower and upper bound.
+
+    #### Return modes
+
+    Succeeds if the value's type is compatible with the target type, and
+    produces a silenceable failure otherwise. Produces a definite failure
+    if the operand is not associated with a single payload value.
+  }];
+
+  let arguments = (ins TransformValueHandleTypeInterface:$operand_handle,
+                       I64Attr:$dim,
+                       OptionalAttr<I64Attr>:$lower_bound,
+                       OptionalAttr<I64Attr>:$upper_bound);
+
+  // `<=` isn't a valid literal so use `le` instead.
+  let assemblyFormat = [{
+    $operand_handle `[` $dim `]` `,` (`umin` `=` $lower_bound^):(`none`)?
+    `,` (`umax` `=` $upper_bound^):(`none`)?  attr-dict `:` type($operand_handle)
+  }];
+  let extraClassDeclaration = SingleValueMatcher.extraDeclaration;
+
+  let cppNamespace = "mlir::iree_compiler::IREE::transform_dialect";
+}
+
 def MatchDimIsMultipleOfOp : Op<Transform_Dialect, "iree.match.dim_is_multiple_of",
     [IsolatedFromAbove,
      MatchOpInterface,
      SingleValueMatcher,
      MemoryEffectsOpInterface]> {
   let summary =
-      "Checks if the body of the target op matches the body of the single contained op";
+      "Checks the static size of a dim is divisible by a given value";
   let description = [{
     Checks whether the given dimension given shaped value is a multiple of the
     given size.