[mlir][Interfaces] ValueBoundsOpInterface: Compute constant bounds

[springerm]

Summary

Add a helper function that computes a constant (int64_t) bound. The stopCondition is optional: If none is provided, the traversal continues until a constant bound could be computed.

Depends On: D145806

[springerm]

rebase

[springerm]

rebase

[springerm]

rebase

[dcaballe]

mlir/lib/Interfaces/ValueBoundsOpInterface.cpp

364

nit: .

372

Did you finally add the option to make UB close? (I guess I'll figure it out :))

mlir/test/Dialect/Affine/value-bounds-reification.mlir

42

For cases where constant is not provided and there are multiple bounds, are we returning now the non-constant one?

88

Nit: What about introducing a test.reify_constant_bound instead?

[springerm]

mlir/lib/Interfaces/ValueBoundsOpInterface.cpp

372

Yes, it will be added as part of D146356.

mlir/test/Dialect/Affine/value-bounds-reification.mlir

42

No, we still return the constant one. That's a limitation of the FlatLinearConstraints implementation. (There's a TODO in there.)

The constant attribute here just says that the computed bound must be a constant. If it is not specified, the computed bound may be constant or not constant.

mlir/include/mlir/Interfaces/ValueBoundsOpInterface.h

 
 #include "mlir/Analysis/FlatValueConstraints.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/Value.h"
 #include "mlir/Interfaces/DestinationStyleOpInterface.h"
 #include "llvm/ADT/SetVector.h"
 
+#include <queue>
+
 namespace mlir {
 
 using ValueDimList = SmallVector<std::pair<Value, std::optional<int64_t>>>;
 
 /// A helper class to be used with `ValueBoundsOpInterface`. This class stores a
 /// constraint system and mapping of constrained variables to index-typed
 /// values or dimension sizes of shaped values.
 ///
public:
   /// given value is visited in a worklist-driven manner and the constraint set
   /// is populated according to `ValueBoundsOpInterface` for each visited value.
   static LogicalResult computeBound(AffineMap &resultMap,
                                     ValueDimList &mapOperands, Location loc,
                                     presburger::BoundType type, Value value,
                                     std::optional<int64_t> dim,
                                     function_ref<bool(Value)> stopCondition);
 
+  /// Compute a constant bound for the given index-typed value or shape
+  /// dimension size.
+  ///
+  /// `dim` must be `nullopt` if and only if `value` is index-typed. This
+  /// function traverses the backward slice of the given value in a
+  /// worklist-driven manner until `stopCondition` evaluates to "true". The
+  /// constraint set is populated according to `ValueBoundsOpInterface` for each
+  /// visited value. (No constraints are added for values for which the stop
+  /// condition evaluates to "true".)
+  ///
+  /// The stop condition is optional: If none is specified, the backward slice
+  /// is traversed in a breadth-first manner until a constant bound could be
+  /// computed.
+  static FailureOr<int64_t>
+  computeConstantBound(presburger::BoundType type, Value value,
+                       std::optional<int64_t> dim = std::nullopt,
+                       function_ref<bool(Value)> stopCondition = nullptr);
+
   /// Bound the given index-typed value by the given expression.
   void addBound(presburger::BoundType type, Value value, AffineExpr expr);
 
   /// Bound the given shaped value dimension by the given expression.
   void addBound(presburger::BoundType type, Value value, int64_t dim,
                 AffineExpr expr);
 
   /// Bound the given index-typed value by a constant or SSA value.
protected:
   void projectOut(function_ref<bool(ValueDim)> condition);
 
   /// Mapping of columns to values/shape dimensions.
   SmallVector<ValueDim> positionToValueDim;
   /// Reverse mapping of values/shape dimensions to columns.
   DenseMap<ValueDim, int64_t> valueDimToPosition;
 
   /// Worklist of values/shape dimensions that have not been processed yet.
-  SetVector<int64_t> worklist;
+  std::queue<int64_t> worklist;
 
   /// Constraint system of equalities and inequalities.
   FlatConstraints cstr;
 
   /// Builder for constructing affine expressions.
   Builder builder;
 };
 

mlir/lib/Interfaces/ValueBoundsOpInterface.cpp

assert((valueDimToPosition.find(valueDim) == valueDimToPosition.end()) &&
          "already mapped");
   int64_t pos = isSymbol ? cstr.appendVar(VarKind::Symbol)
                          : cstr.appendVar(VarKind::SetDim);
   positionToValueDim.insert(positionToValueDim.begin() + pos, valueDim);
   // Update reverse mapping.
   for (int64_t i = pos; i < positionToValueDim.size(); ++i)
     valueDimToPosition[positionToValueDim[i]] = i;
 
-  worklist.insert(pos);
+  worklist.push(pos);
   return pos;
 }
 
 int64_t ValueBoundsConstraintSet::getPos(Value value, int64_t dim) const {
 #ifndef NDEBUG
   assertValidValueDim(value, dim == kIndexValue
                                  ? std::nullopt
                                  : std::make_optional<int64_t>(dim));
 #endif // NDEBUG
 
   auto it = valueDimToPosition.find(std::make_pair(value, dim));
   assert(it != valueDimToPosition.end() && "expected mapped entry");
   return it->second;
 }
 
 static Operation *getOwnerOfValue(Value value) {
   if (auto bbArg = value.dyn_cast<BlockArgument>())
     return bbArg.getOwner()->getParentOp();
   return value.getDefiningOp();
 }
 
 void ValueBoundsConstraintSet::processWorklist(
     function_ref<bool(Value)> stopCondition) {
   while (!worklist.empty()) {
-    int64_t pos = worklist.pop_back_val();
+    int64_t pos = worklist.front();
+    worklist.pop();
     ValueDim valueDim = positionToValueDim[pos];
     Value value = valueDim.first;
     int64_t dim = valueDim.second;
 
     // Check for static dim size.
     if (dim != kIndexValue) {
       auto shapedType = value.getType().cast<ShapedType>();
       if (shapedType.hasRank() && !shapedType.isDynamicDim(dim)) {
while (nextPos < positionToValueDim.size()) {
     }
   }
 }
 
 LogicalResult ValueBoundsConstraintSet::computeBound(
     AffineMap &resultMap, ValueDimList &mapOperands, Location loc,
     presburger::BoundType type, Value value, std::optional<int64_t> dim,
     function_ref<bool(Value)> stopCondition) {
+#ifndef NDEBUG
+  assertValidValueDim(value, dim);
+#endif // NDEBUG
+
   // Process the backward slice of `value` (i.e., reverse use-def chain) until
   // `stopCondition` is met.
   ValueDim valueDim = std::make_pair(value, dim.value_or(kIndexValue));
   ValueBoundsConstraintSet cstr(valueDim);
   cstr.processWorklist(stopCondition);
 
   // Project out all variables (apart from `valueDim`) that do not match the
   // stop condition.
assert(value.getType().cast<ShapedType>().isDynamicDim(dim) &&
            "expected dynamic dim");
     mapOperands.push_back(std::make_pair(value, dim));
   }
 
   resultMap = bound.replaceDimsAndSymbols(replacementDims, replacementSymbols,
                                           numDims, numSymbols);
   return success();
 }
+
+FailureOr<int64_t> ValueBoundsConstraintSet::computeConstantBound(
+    presburger::BoundType type, Value value, std::optional<int64_t> dim,
+    function_ref<bool(Value)> stopCondition) {
+#ifndef NDEBUG
+  assertValidValueDim(value, dim);
+#endif // NDEBUG
+
+  // Process the backward slice of `value` (i.e., reverse use-def chain) until
+  // `stopCondition` is met.
+  ValueDim valueDim = std::make_pair(value, dim.value_or(kIndexValue));
+  ValueBoundsConstraintSet cstr(valueDim);
+  int64_t pos = cstr.getPos(value, dim.value_or(kIndexValue));
+  if (stopCondition) {
+    cstr.processWorklist(stopCondition);
+  } else {
+    // No stop condition specified: Keep adding constraints until a bound could
+    // be computed
+    cstr.processWorklist(/*stopCondition=*/[&](Value v) {
+      return cstr.cstr.getConstantBound64(type, pos).has_value();
+    });
+  }
+
+  // Compute constant bound for `valueDim`.
+  if (auto bound = cstr.cstr.getConstantBound64(type, pos))
+    return type == BoundType::UB ? *bound + 1 : *bound;
+  return failure();
+}

mlir/test/Dialect/Affine/value-bounds-reification.mlir

func.func @reify_through_chain(%sz0: index, %sz2: index) -> (index, index, index) {
   return %4, %5, %6 : index, index, index
 }
 
 // -----
 
 // CHECK-LABEL: func @reify_slice_bound(
 //       CHECK:   %[[c5:.*]] = arith.constant 5 : index
 //       CHECK:   "test.some_use"(%[[c5]])
+//       CHECK:   %[[c5:.*]] = arith.constant 5 : index
+//       CHECK:   "test.some_use"(%[[c5]])
 func.func @reify_slice_bound(%t: tensor<?x?xi32>, %idx: index, %ub: index, %f: f32) {
   %c0 = arith.constant 0 : index
   %c4 = arith.constant 4 : index
   scf.for %iv = %c0 to %ub step %c4 {
     %sz = affine.min affine_map<(d0)[s0] -> (-d0 + s0, 4)>(%iv)[%ub]
     %slice = tensor.extract_slice %t[%idx, %iv] [1, %sz] [1, 1] : tensor<?x?xi32> to tensor<1x?xi32>
     %filled = linalg.fill ins(%f : f32) outs(%slice : tensor<1x?xi32>) -> tensor<1x?xi32>
+
     %bound = "test.reify_bound"(%filled) {dim = 1, type = "UB"} : (tensor<1x?xi32>) -> (index)
     "test.some_use"(%bound) : (index) -> ()
+
+    %bound_const = "test.reify_bound"(%filled) {dim = 1, type = "UB", constant} : (tensor<1x?xi32>) -> (index)
+    "test.some_use"(%bound_const) : (index) -> ()
   }
   return
 }
 
 // -----
 
 // CHECK: #[[$map:.*]] = affine_map<()[s0, s1] -> (s0 - s1 + 1)>
 // CHECK-LABEL: func @scf_for(
scf.for %iv0 = %lb0 to %ub0 step %step0 {
     "test.some_use"(%ub1_ub) : (index) -> ()
 
     // CHECK: %[[c129:.*]] = arith.constant 129 : index
     // CHECK: "test.some_use"(%[[c129]])
     %lb1 = affine.apply affine_map<()[s0] -> ((s0 floordiv 32) * 32)>()[%ub1]
     %lb1_ub = "test.reify_bound"(%lb1) {type = "UB"} : (index) -> (index)
     "test.some_use"(%lb1_ub) : (index) -> ()
 
+    // CHECK: %[[c129:.*]] = arith.constant 129 : index
+    // CHECK: "test.some_use"(%[[c129]])
+    %lb1_ub_const = "test.reify_bound"(%lb1) {type = "UB", constant} : (index) -> (index)
+    "test.some_use"(%lb1_ub_const) : (index) -> ()
+
     scf.for %iv1 = %lb1 to %ub1 step %c32 {
       // CHECK: %[[c32:.*]] = arith.constant 32 : index
       // CHECK: "test.some_use"(%[[c32]])
       %sz = affine.apply affine_map<(d0)[s0] -> (-d0 + s0)>(%iv1)[%ub1]
       %sz_ub = "test.reify_bound"(%sz) {type = "UB"} : (index) -> (index)
       "test.some_use"(%sz_ub) : (index) -> ()
 
       scf.for %iv2 = %c0 to %ub2 step %c1 {
         %slice1 = tensor.extract_slice %t1[0, %iv2] [1, 1] [1, 1] : tensor<1x?xi8> to tensor<1x1xi8>
         %slice2 = tensor.extract_slice %t2[%iv2, 0] [1, %sz] [1, 1] : tensor<?x?xi8> to tensor<1x?xi8>
         %slice3 = tensor.extract_slice %t3[0, 0] [1, %sz] [1, 1] : tensor<1x?xi32> to tensor<1x?xi32>
         %matmul = linalg.matmul ins(%slice1, %slice2 : tensor<1x1xi8>, tensor<1x?xi8>) outs(%slice3 : tensor<1x?xi32>) -> tensor<1x?xi32>
 
         // CHECK: %[[c32:.*]] = arith.constant 32 : index
         // CHECK: "test.some_use"(%[[c32]])
         %matmul_ub = "test.reify_bound"(%matmul) {dim = 1, type = "UB"} : (tensor<1x?xi32>) -> (index)
         "test.some_use"(%matmul_ub) : (index) -> ()
+
+        // CHECK: %[[c32:.*]] = arith.constant 32 : index
+        // CHECK: "test.some_use"(%[[c32]])
+        %matmul_ub_const = "test.reify_bound"(%matmul) {dim = 1, type = "UB", constant} : (tensor<1x?xi32>) -> (index)
+        "test.some_use"(%matmul_ub_const) : (index) -> ()
       }
     }
   }
   return
 }

mlir/test/Dialect/Tensor/value-bounds-op-interface-impl.mlir

 func.func @extract_slice_static(%t: tensor<?xf32>) -> index {
   %0 = tensor.extract_slice %t[2][5][1] : tensor<?xf32> to tensor<5xf32>
   %1 = "test.reify_bound"(%0) {dim = 0} : (tensor<5xf32>) -> (index)
   return %1 : index
 }
 
 // -----
 
+func.func @extract_slice_dynamic_constant(%t: tensor<?xf32>, %sz: index) -> index {
+  %0 = tensor.extract_slice %t[2][%sz][1] : tensor<?xf32> to tensor<?xf32>
+  // expected-error @below{{could not reify bound}}
+  %1 = "test.reify_bound"(%0) {dim = 0, constant} : (tensor<?xf32>) -> (index)
+  return %1 : index
+}
+
+// -----
+
+// CHECK-LABEL: func @extract_slice_static_constant(
+//  CHECK-SAME:     %[[t:.*]]: tensor<?xf32>
+//       CHECK:   %[[c5:.*]] = arith.constant 5 : index
+//       CHECK:   return %[[c5]]
+func.func @extract_slice_static_constant(%t: tensor<?xf32>) -> index {
+  %0 = tensor.extract_slice %t[2][5][1] : tensor<?xf32> to tensor<5xf32>
+  %1 = "test.reify_bound"(%0) {dim = 0, constant} : (tensor<5xf32>) -> (index)
+  return %1 : index
+}
+
+// -----
+
 // CHECK-LABEL: func @extract_slice_rank_reduce(
 //  CHECK-SAME:     %[[t:.*]]: tensor<?x?xf32>, %[[sz:.*]]: index
 //       CHECK:   return %[[sz]]
 func.func @extract_slice_rank_reduce(%t: tensor<?x?xf32>, %sz: index) -> index {
   %0 = tensor.extract_slice %t[0, 2][1, %sz][1, 1] : tensor<?x?xf32> to tensor<?xf32>
   %1 = "test.reify_bound"(%0) {dim = 0} : (tensor<?xf32>) -> (index)
   return %1 : index
 }

mlir/test/lib/Dialect/Affine/TestReifyValueBounds.cpp

if (op->getName().getStringRef() == "test.reify_bound") {
       }
 
       // Get shape dimension (if any).
       auto dim = value.getType().isIndex()
                      ? std::nullopt
                      : std::make_optional<int64_t>(
                            op->getAttrOfType<IntegerAttr>("dim").getInt());
 
-      // Reify value bound.
-      rewriter.setInsertionPointAfter(op);
-      FailureOr<OpFoldResult> reified;
-      if (!reifyToFuncArgs) {
-        // Reify in terms of the op's operands.
-        reified =
-            reifyValueBound(rewriter, op->getLoc(), *boundType, value, dim);
-      } else {
+      // Check if a constant was requested.
+      bool constant = op->hasAttr("constant");
+
+      // Prepare stop condition. By default, reify in terms of the op's
+      // operands. No stop condition is used when a constant was requested.
+      std::function<bool(Value)> stopCondition = [&](Value v) {
+        // Reify in terms of SSA values that are different from `value`.
+        return v != value;
+      };
+      if (reifyToFuncArgs) {
         // Reify in terms of function block arguments.
-        auto stopCondition = [](Value v) {
+        stopCondition = stopCondition = [](Value v) {
           auto bbArg = v.dyn_cast<BlockArgument>();
           if (!bbArg)
             return false;
           return isa<FunctionOpInterface>(
               bbArg.getParentBlock()->getParentOp());
         };
+      }
+
+      // Reify value bound
+      rewriter.setInsertionPointAfter(op);
+      FailureOr<OpFoldResult> reified = failure();
+      if (constant) {
+        auto reifiedConst = ValueBoundsConstraintSet::computeConstantBound(
+            *boundType, value, dim, /*stopCondition=*/nullptr);
+        if (succeeded(reifiedConst))
+          reified =
+              FailureOr<OpFoldResult>(rewriter.getIndexAttr(*reifiedConst));
+      } else {
         reified = reifyValueBound(rewriter, op->getLoc(), *boundType, value,
                                   dim, stopCondition);
       }
       if (failed(reified)) {
         op->emitOpError("could not reify bound");
         return WalkResult::interrupt();
       }