[Preprocessing] Use value bounds for dim multiple checking (#20583)

This also adds support for udiv to util.assume value bounds constraints.

In the future we can add more constraints to represent udiv, but for now this
implements at least one working option.

Author: qedawkins

compiler/src/iree/compiler/ExternalInterfaces/UtilExternalModels.cpp

@@ -126,12 +126,30 @@ struct UtilAssumeIntValueBoundsOpInterface
     auto [min, max] =
         assumeOp.getUnionedUnsignedRange(result.getResultNumber());
 
+    std::optional<int64_t> udiv =
+        assumeOp.getUnionedUnsignedDivisor(result.getResultNumber());
+
     if (min) {
       cstr.bound(result) >= *min;
     }
     if (max) {
       cstr.bound(result) <= *max;
     }
+    if (udiv) {
+      // To represent the divisibility guarantee, emit a bound clamping the
+      // value to the udiv value. i.e.
+      //
+      // v == floordiv(v, udiv) * udiv
+      //
+      // Mod/divide folders can cleanup such terms with the appropriate bounds
+      // query.
+      AffineExpr expr =
+          cstr.getExpr(assumeOp.getOperand(result.getResultNumber()));
+      AffineExpr udivCst =
+          getAffineConstantExpr(udiv.value(), op->getContext());
+      AffineExpr clampExpr = expr.floorDiv(udivCst) * udivCst;
+      cstr.bound(result) == clampExpr;
+    }
   }
 };
 

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

@@ -149,7 +149,7 @@ 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
+%0 = util.assume.int %arg0<umin = 12, umax = 16, udiv = 4> : index
 %input = tensor.empty(%0) : tensor<?xf32>
 //       CHECK: call @external
 //  CHECK-SAME:   match_status = "both_matched"
@@ -178,6 +178,7 @@ module attributes {transform.with_named_sequence} {
     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
+    transform.iree.match.dim_is_multiple_of %in0[0], 2 : !transform.any_value
     %0 = transform.param.constant "both_matched" -> !transform.any_param
     transform.yield %call, %0 : !transform.any_op, !transform.any_param
   }

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

@@ -327,6 +327,7 @@ DiagnosedSilenceableFailure
 IREE::transform_dialect::MatchDimIsMultipleOfOp::matchValue(
     Value current, transform::TransformResults &results,
     transform::TransformState &state) {
+  MLIRContext *ctx = current.getContext();
   auto shapedType = dyn_cast<ShapedType>(current.getType());
   if (!shapedType) {
     return emitSilenceableError()
@@ -338,7 +339,20 @@ IREE::transform_dialect::MatchDimIsMultipleOfOp::matchValue(
            << "dim " << dim << " out of range for shaped type " << shapedType;
   }
   int64_t size = getSize();
-  if (shapedType.getShape()[dim] % size != 0) {
+  ValueBoundsConstraintSet::Variable dimVar(current, dim);
+
+  // Check if current[dim] % size == 0. There are a couple of options for how
+  // to do this (e.g. mul(floordiv)). Affine map canonicalizations are good
+  // at dropping terms that statically divide the mod RHS so we go with this
+  // one.
+  AffineMap modMap = AffineMap::get(/*dimCount=*/0, /*symbolCount=*/1,
+                                    getAffineSymbolExpr(0, ctx) %
+                                        getAffineConstantExpr(size, ctx));
+  ValueBoundsConstraintSet::Variable modVar(modMap, {dimVar});
+  Builder b(ctx);
+  FailureOr<bool> maybeFailed = ValueBoundsConstraintSet::areEqual(
+      modVar, OpFoldResult{b.getIndexAttr(0)});
+  if (failed(maybeFailed) || !maybeFailed.value()) {
     return emitSilenceableError()
            << "dim " << dim << " of shaped type " << shapedType
            << " is not a multiple of " << size;

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

@@ -118,6 +118,7 @@ def MatchDimBoundsOp : Op<Transform_Dialect, "iree.match.dim_bounds",
   let cppNamespace = "mlir::iree_compiler::IREE::transform_dialect";
 }
 
+// TODO: Combine with MatchDimBounds.
 def MatchDimIsMultipleOfOp : Op<Transform_Dialect, "iree.match.dim_is_multiple_of",
     [IsolatedFromAbove,
      MatchOpInterface,