[mlir][SCF] ValueBoundsOpInterface: Support `scf.for` results and iter_args

[springerm]

Summary

If an scf.for loop yields an equal index-typed value or a shaped value with the same dimension sizes (in comparison to the corresponding iter_arg), bounds can be computed for the iter_arg and the OpResult of the scf.for op.

This change also adds a dim parameter to stopCondition.

Depends On: D146296

[springerm]

rebase

[springerm]

update API

[dcaballe]

LGTM! Minor comments

mlir/include/mlir/Interfaces/ValueBoundsOpInterface.h

40–41

doc about the arguments?

73

Could you elaborate a bit more on that dependencies are actually used to create the stopCondition function?

mlir/lib/Dialect/Affine/Transforms/ReifyValueBounds.cpp

21–22

update doc

mlir/lib/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.cpp

26

It looks like the name dim is no longer accurate? Fof scf.for it represents the result index, which is not really a dim? I can't think of a better name, though, other than index...

30

-> dyn_cast new convention

39

I guess there is no gracefully way fail here because if we move forward we might return a incorrect bound...

66

dyn_cast

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

46

This is pretty cool! We are able to infer that we are actually returning %arg, which is %i

[springerm]

mlir/lib/Dialect/Affine/Transforms/ReifyValueBounds.cpp

21–22

Nothing to update here. The new parameter is not used. But I fixed a typo.

mlir/lib/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.cpp

26

dim is still the right name here. This does not refer to the induction variable here. E.g., we could be trying to compute a bound for a shape typed iter_arg.

scf.for ... iter_args(%arg0 = %t) -> tensor<?xf32> {
  "test.reify_bound"(%arg0) {dim = 0}
}

30

This does not work for OpFoldResult yet.

39

The region of the op is SizedRegion<1>:$region, so it guaranteed to have only 1 block. This is just for extra safety.

66

Does not work for OpFoldResult yet.

mlir/include/mlir/Dialect/Affine/Transforms/Transforms.h

 FailureOr<OpFoldResult> reifyValueBound(OpBuilder &b, Location loc,
                                         presburger::BoundType type, Value value,
                                         std::optional<int64_t> dim);
 
 /// Reify a bound for the given index-typed value or shape dimension size in
 /// terms of SSA values for which `stopCondition` is met. `dim` must be
 /// `nullopt` if and only if `value` is index-typed. LB and EQ bounds are
 /// closed, UB bounds are open.
-FailureOr<OpFoldResult>
-reifyValueBound(OpBuilder &b, Location loc, presburger::BoundType type,
-                Value value, std::optional<int64_t> dim,
-                function_ref<bool(Value)> stopCondition);
+FailureOr<OpFoldResult> reifyValueBound(
+    OpBuilder &b, Location loc, presburger::BoundType type, Value value,
+    std::optional<int64_t> dim,
+    function_ref<bool(Value, std::optional<int64_t>)> stopCondition);
 
 } // namespace mlir
 
 #endif // MLIR_DIALECT_AFFINE_TRANSFORMS_TRANSFORMS_H

mlir/include/mlir/Interfaces/ValueBoundsOpInterface.h

 /// expressed in terms of another value/dimension, `getExpr` can be used to
 /// retrieve an AffineExpr that represents the specified value/dimension.
 ///
 /// When a value/dimension is retrieved for the first time through `getExpr`,
 /// it is added to an internal worklist. See `computeBound` for more details.
 ///
 /// Note: Any modification of the IR invalides the data stored in this class.
 class ValueBoundsConstraintSet {
 public:
   /// Compute a bound for the given index-typed value or shape dimension size.
   /// The computed bound is stored in `resultMap`. The operands of the bound are
   /// stored in `mapOperands`. An operand is either an index-type SSA value
   /// or a shaped value and a dimension.
   ///
   /// `dim` must be `nullopt` if and only if `value` is index-typed. The bound
-  /// is computed in terms of values for which `stopCondition` evaluates to
-  /// "true". To that end, the backward slice (reverse use-def chain) of the
-  /// given value is visited in a worklist-driven manner and the constraint set
-  /// is populated according to `ValueBoundsOpInterface` for each visited value.
+  /// is computed in terms of values/dimensions for which `stopCondition`
+  /// evaluates to "true". To that end, the backward slice (reverse use-def
+  /// chain) of the given value is visited in a worklist-driven manner and the
+  /// constraint set is populated according to `ValueBoundsOpInterface` for each
+  /// visited value.
+  ///
+  /// Note: `stopCondition` should not evaluate to "true" for the input
+  /// value/dimension.
+  static LogicalResult
+  computeBound(AffineMap &resultMap, ValueDimList &mapOperands,
+               presburger::BoundType type, Value value,
+               std::optional<int64_t> dim,
+               function_ref<bool(Value, std::optional<int64_t>)> stopCondition);
+
+  /// Compute a bound in terms of the values/dimensions in `dependencies`.
   static LogicalResult computeBound(AffineMap &resultMap,
                                     ValueDimList &mapOperands,
                                     presburger::BoundType type, Value value,
                                     std::optional<int64_t> dim,
-                                    function_ref<bool(Value)> stopCondition);
+                                    ValueDimList dependencies);
 
   /// 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);
+  static FailureOr<int64_t> computeConstantBound(
+      presburger::BoundType type, Value value,
+      function_ref<bool(Value, std::optional<int64_t>)> 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);
 
protected:
   static constexpr int64_t kIndexValue = -1;
 
   using ValueDim = std::pair<Value, int64_t>;
   ValueBoundsConstraintSet(ValueDim valueDim);
 
   /// Iteratively process all elements on the worklist until an index-typed
   /// value or shaped value meets `stopCondition`. Such values are not processed
   /// any further.
-  void processWorklist(function_ref<bool(Value)> stopCondition);
+  void processWorklist(
+      function_ref<bool(Value, std::optional<int64_t>)> stopCondition);
 
   /// Bound the given column in the underlying constraint set by the given
   /// expression.
   void addBound(presburger::BoundType type, int64_t pos, AffineExpr expr);
 
   /// Return the column position of the given value/dimension. Asserts that the
   /// value/dimension exists in the constraint set.
   int64_t getPos(Value value, int64_t dim = kIndexValue) const;

mlir/lib/Dialect/Affine/Transforms/ReifyValueBounds.cpp

 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Interfaces/ValueBoundsOpInterface.h"
 
 using namespace mlir;
 
 FailureOr<OpFoldResult> mlir::reifyValueBound(OpBuilder &b, Location loc,
                                               presburger::BoundType type,
                                               Value value,
                                               std::optional<int64_t> dim) {
-  auto stopCondition = [&](Value v) {
+  auto stopCondition = [&](Value v, std::optional<int64_t> d) {
     // Reify in terms of SSA values that are different from `value`.
     return v != value;
   };
   return reifyValueBound(b, loc, type, value, dim, stopCondition);
 }
 
-FailureOr<OpFoldResult>
-mlir::reifyValueBound(OpBuilder &b, Location loc, presburger::BoundType type,
-                      Value value, std::optional<int64_t> dim,
-                      function_ref<bool(Value)> stopCondition) {
+FailureOr<OpFoldResult> mlir::reifyValueBound(
+    OpBuilder &b, Location loc, presburger::BoundType type, Value value,
+    std::optional<int64_t> dim,
+    function_ref<bool(Value, std::optional<int64_t>)> stopCondition) {
   // Compute bound.
   AffineMap boundMap;
   ValueDimList mapOperands;
   if (failed(ValueBoundsConstraintSet::computeBound(boundMap, mapOperands, type,
                                                     value, dim, stopCondition)))
     return failure();
 
   // Materialize tensor.dim/memref.dim ops.

mlir/lib/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.cpp

 namespace scf {
 namespace {
 
 struct ForOpInterface
     : public ValueBoundsOpInterface::ExternalModel<ForOpInterface, ForOp> {
   void populateBoundsForIndexValue(Operation *op, Value value,
                                    ValueBoundsConstraintSet &cstr) const {
     auto forOp = cast<ForOp>(op);
-    // Only IV is supported at the moment.
-    if (value != forOp.getInductionVar())
-      return;
 
+    if (value == forOp.getInductionVar()) {
       // TODO: Take into account step size.
       cstr.addBound(BoundType::LB, value, cstr.getExpr(forOp.getLowerBound()));
       cstr.addBound(BoundType::UB, value, cstr.getExpr(forOp.getUpperBound()));
+      return;
+    }
+
+    // `value` is an iter_arg or an OpResult.
+    int64_t iterArgIdx;
+    if (auto iterArg = value.dyn_cast<BlockArgument>()) {
+      iterArgIdx = iterArg.getArgNumber() - forOp.getNumInductionVars();
+    } else {
+      iterArgIdx = value.cast<OpResult>().getResultNumber();
+    }
+
+    // An EQ constraint can be added if the yielded value is the same value as
+    // the corresponding block argument.
+    assert(forOp.getLoopBody().hasOneBlock() &&
+           "multiple blocks not supported");
+    Value yieldedValue =
+        cast<scf::YieldOp>(forOp.getLoopBody().front().getTerminator())
+            .getOperand(iterArgIdx);
+    Value iterArg = forOp.getRegionIterArg(iterArgIdx);
+    Value initArg = forOp.getInitArgs()[iterArgIdx];
+
+    // Compute EQ bound for yielded value.
+    AffineMap bound;
+    ValueDimList boundOperands;
+    LogicalResult status = ValueBoundsConstraintSet::computeBound(
+        bound, boundOperands, BoundType::EQ, yieldedValue,
+        /*dim=*/std::nullopt, {{iterArg, std::nullopt}});
+    if (failed(status))
+      return;
+    if (bound.getNumResults() != 1)
+      return;
+
+    // Check if computed bound equals the corresponding iter_arg.
+    Value singleValue = nullptr;
+    if (auto dimExpr = bound.getResult(0).dyn_cast<AffineDimExpr>()) {
+      int64_t idx = dimExpr.getPosition();
+      if (!boundOperands[idx].second.has_value())
+        singleValue = boundOperands[idx].first;
+    } else if (auto symExpr = bound.getResult(0).dyn_cast<AffineSymbolExpr>()) {
+      int64_t idx = symExpr.getPosition() + bound.getNumDims();
+      if (!boundOperands[idx].second.has_value())
+        singleValue = boundOperands[idx].first;
+    }
+    if (singleValue == iterArg)
+      cstr.addBound(BoundType::EQ, value, cstr.getExpr(initArg));
   }
 
   void populateBoundsForShapedValueDim(Operation *op, Value value, int64_t dim,
                                        ValueBoundsConstraintSet &cstr) const {
-    // iter_arg / return value not supported.
+    auto forOp = cast<ForOp>(op);
+
+    // `value` is an iter_arg or an OpResult.
+    int64_t iterArgIdx;
+    if (auto iterArg = value.dyn_cast<BlockArgument>()) {
+      iterArgIdx = iterArg.getArgNumber() - forOp.getNumInductionVars();
+    } else {
+      iterArgIdx = value.cast<OpResult>().getResultNumber();
+    }
+
+    // An EQ constraint can be added if the dimension size of the yielded value
+    // equals the dimensions size of the corresponding block argument.
+    assert(forOp.getLoopBody().hasOneBlock() &&
+           "multiple blocks not supported");
+    Value yieldedValue =
+        cast<scf::YieldOp>(forOp.getLoopBody().front().getTerminator())
+            .getOperand(iterArgIdx);
+    Value iterArg = forOp.getRegionIterArg(iterArgIdx);
+    Value initArg = forOp.getInitArgs()[iterArgIdx];
+
+    // Compute EQ bound for yielded value.
+    AffineMap bound;
+    ValueDimList boundOperands;
+    LogicalResult status = ValueBoundsConstraintSet::computeBound(
+        bound, boundOperands, BoundType::EQ, yieldedValue, dim,
+        {{iterArg, dim}});
+    if (failed(status))
+      return;
+    if (bound.getNumResults() != 1)
       return;
+
+    // Check if computed bound equals the corresponding iter_arg.
+    Value singleValue = nullptr;
+    int64_t singleDim = -1;
+    if (auto dimExpr = bound.getResult(0).dyn_cast<AffineDimExpr>()) {
+      int64_t idx = dimExpr.getPosition();
+      if (boundOperands[idx].second.has_value()) {
+        singleValue = boundOperands[idx].first;
+        singleDim = *boundOperands[idx].second;
+      }
+    } else if (auto symExpr = bound.getResult(0).dyn_cast<AffineSymbolExpr>()) {
+      int64_t idx = symExpr.getPosition() + bound.getNumDims();
+      if (boundOperands[idx].second.has_value()) {
+        singleValue = boundOperands[idx].first;
+        singleDim = *boundOperands[idx].second;
+      }
+    }
+    if (singleValue == iterArg && singleDim == dim)
+      cstr.addBound(BoundType::EQ, value, dim, cstr.getExpr(initArg, dim));
   }
 };
 
 } // namespace
 } // namespace scf
 } // namespace mlir
 
 void mlir::scf::registerValueBoundsOpInterfaceExternalModels(
     DialectRegistry &registry) {
   registry.addExtension(+[](MLIRContext *ctx, scf::SCFDialect *dialect) {
     scf::ForOp::attachInterface<scf::ForOpInterface>(*ctx);
   });
 }

mlir/lib/Interfaces/ValueBoundsOpInterface.cpp

 
 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) {
+    function_ref<bool(Value, std::optional<int64_t>)> stopCondition) {
   while (!worklist.empty()) {
     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)) {
         addBound(BoundType::EQ, value, dim,
                  builder.getAffineConstantExpr(shapedType.getDimSize(dim)));
         continue;
       }
     }
 
     // Do not process any further if the stop condition is met.
-    if (stopCondition(value))
+    auto maybeDim = dim == kIndexValue ? std::nullopt : std::make_optional(dim);
+    if (stopCondition(value, maybeDim))
       continue;
 
     // Query `ValueBoundsOpInterface` for constraints. New items may be added to
     // the worklist.
     auto valueBoundsOp =
         dyn_cast<ValueBoundsOpInterface>(getOwnerOfValue(value));
     if (!valueBoundsOp)
       continue;
if (condition(positionToValueDim[nextPos])) {
       ++nextPos;
     }
   }
 }
 
 LogicalResult ValueBoundsConstraintSet::computeBound(
     AffineMap &resultMap, ValueDimList &mapOperands, presburger::BoundType type,
     Value value, std::optional<int64_t> dim,
-    function_ref<bool(Value)> stopCondition) {
+    function_ref<bool(Value, std::optional<int64_t>)> stopCondition) {
 #ifndef NDEBUG
   assertValidValueDim(value, dim);
+  assert(!stopCondition(value, dim) &&
+         "stop condition should not be satisfied for starting point");
 #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.
   cstr.projectOut([&](ValueDim p) {
     // Do not project out `valueDim`.
     if (valueDim == p)
       return false;
-    return !stopCondition(p.first);
+    auto maybeDim =
+        p.second == kIndexValue ? std::nullopt : std::make_optional(p.second);
+    return !stopCondition(p.first, maybeDim);
   });
 
   // Compute lower and upper bounds for `valueDim`.
   int64_t pos = cstr.getPos(value, dim.value_or(kIndexValue));
   SmallVector<AffineMap> lb(1), ub(1);
   cstr.cstr.getSliceBounds(pos, 1, value.getContext(), &lb, &ub,
                            /*getClosedUB=*/true);
 
for (int64_t i = 0; i < cstr.cstr.getNumDimAndSymbolVars(); ++i) {
     mapOperands.push_back(std::make_pair(value, dim));
   }
 
   resultMap = bound.replaceDimsAndSymbols(replacementDims, replacementSymbols,
                                           numDims, numSymbols);
   return success();
 }
 
+LogicalResult ValueBoundsConstraintSet::computeBound(
+    AffineMap &resultMap, ValueDimList &mapOperands, presburger::BoundType type,
+    Value value, std::optional<int64_t> dim, ValueDimList dependencies) {
+  if (llvm::is_contained(dependencies, std::make_pair(value, dim))) {
+    resultMap = AffineMap::get(/*dimCount=*/1, /*symbolCount=*/0,
+                               Builder(value.getContext()).getAffineDimExpr(0));
+    mapOperands.clear();
+    mapOperands.push_back(std::make_pair(value, dim));
+    return success();
+  }
+
+  return computeBound(resultMap, mapOperands, type, value, dim,
+                      [&](Value v, std::optional<int64_t> d) {
+                        return llvm::is_contained(dependencies,
+                                                  std::make_pair(v, d));
+                      });
+}
+
 FailureOr<int64_t> ValueBoundsConstraintSet::computeConstantBound(
     presburger::BoundType type, Value value, std::optional<int64_t> dim,
-    function_ref<bool(Value)> stopCondition) {
+    function_ref<bool(Value, std::optional<int64_t>)> 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) {
+    cstr.processWorklist(
+        /*stopCondition=*/[&](Value v, std::optional<int64_t> dim) {
           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/SCF/value-bounds-op-interface-impl.mlir

 // RUN: mlir-opt %s -test-affine-reify-value-bounds -verify-diagnostics \
 // RUN:     -split-input-file | FileCheck %s
 
 // CHECK-LABEL: func @scf_for(
 //  CHECK-SAME:     %[[a:.*]]: index, %[[b:.*]]: index, %[[c:.*]]: index
 //       CHECK:   "test.some_use"(%[[a]], %[[b]])
 func.func @scf_for(%a: index, %b: index, %c: index) {
   scf.for %iv = %a to %b step %c {
     %0 = "test.reify_bound"(%iv) {type = "LB"} : (index) -> (index)
     %1 = "test.reify_bound"(%iv) {type = "UB"} : (index) -> (index)
     "test.some_use"(%0, %1) : (index, index) -> ()
   }
   return
 }
+
+// -----
+
+// CHECK-LABEL: func @scf_for_index_result_small(
+//  CHECK-SAME:     %[[i:.*]]: index, %[[a:.*]]: index, %[[b:.*]]: index, %[[c:.*]]: index
+//       CHECK:   "test.some_use"(%[[i]])
+//       CHECK:   "test.some_use"(%[[i]])
+func.func @scf_for_index_result_small(%i: index, %a: index, %b: index, %c: index) {
+  %0 = scf.for %iv = %a to %b step %c iter_args(%arg = %i) -> index {
+    %1 = "test.reify_bound"(%arg) {type = "EQ"} : (index) -> (index)
+    "test.some_use"(%1) : (index) -> ()
+    scf.yield %arg : index
+  }
+  %2 = "test.reify_bound"(%0) {type = "EQ"} : (index) -> (index)
+  "test.some_use"(%2) : (index) -> ()
+  return
+}
+
+// -----
+
+// CHECK-LABEL: func @scf_for_index_result(
+//  CHECK-SAME:     %[[i:.*]]: index, %[[a:.*]]: index, %[[b:.*]]: index, %[[c:.*]]: index
+//       CHECK:   "test.some_use"(%[[i]])
+//       CHECK:   "test.some_use"(%[[i]])
+func.func @scf_for_index_result(%i: index, %a: index, %b: index, %c: index) {
+  %0 = scf.for %iv = %a to %b step %c iter_args(%arg = %i) -> index {
+    %add = arith.addi %arg, %a : index
+    %sub = arith.subi %add, %a : index
+
+    %1 = "test.reify_bound"(%arg) {type = "EQ"} : (index) -> (index)
+    "test.some_use"(%1) : (index) -> ()
+    scf.yield %sub : index
+  }
+  %2 = "test.reify_bound"(%0) {type = "EQ"} : (index) -> (index)
+  "test.some_use"(%2) : (index) -> ()
+  return
+}
+
+// -----
+
+// CHECK-LABEL: func @scf_for_tensor_result_small(
+//  CHECK-SAME:     %[[t:.*]]: tensor<?xf32>, %[[a:.*]]: index, %[[b:.*]]: index, %[[c:.*]]: index
+//       CHECK:   %[[dim:.*]] = tensor.dim %[[t]]
+//       CHECK:   "test.some_use"(%[[dim]])
+//       CHECK:   %[[dim:.*]] = tensor.dim %[[t]]
+//       CHECK:   "test.some_use"(%[[dim]])
+func.func @scf_for_tensor_result_small(%t: tensor<?xf32>, %a: index, %b: index, %c: index) {
+  %0 = scf.for %iv = %a to %b step %c iter_args(%arg = %t) -> tensor<?xf32> {
+    %1 = "test.reify_bound"(%arg) {type = "EQ", dim = 0} : (tensor<?xf32>) -> (index)
+    "test.some_use"(%1) : (index) -> ()
+    scf.yield %arg : tensor<?xf32>
+  }
+  %2 = "test.reify_bound"(%0) {type = "EQ", dim = 0} : (tensor<?xf32>) -> (index)
+  "test.some_use"(%2) : (index) -> ()
+  return
+}
+
+// -----
+
+// CHECK-LABEL: func @scf_for_tensor_result(
+//  CHECK-SAME:     %[[t:.*]]: tensor<?xf32>, %[[a:.*]]: index, %[[b:.*]]: index, %[[c:.*]]: index
+//       CHECK:   %[[dim:.*]] = tensor.dim %[[t]]
+//       CHECK:   "test.some_use"(%[[dim]])
+//       CHECK:   %[[dim:.*]] = tensor.dim %[[t]]
+//       CHECK:   "test.some_use"(%[[dim]])
+func.func @scf_for_tensor_result(%t: tensor<?xf32>, %a: index, %b: index, %c: index) {
+  %cst = arith.constant 5.0 : f32
+  %0 = scf.for %iv = %a to %b step %c iter_args(%arg = %t) -> tensor<?xf32> {
+    %filled = linalg.fill ins(%cst : f32) outs(%arg : tensor<?xf32>) -> tensor<?xf32>
+    %1 = "test.reify_bound"(%arg) {type = "EQ", dim = 0} : (tensor<?xf32>) -> (index)
+    "test.some_use"(%1) : (index) -> ()
+    scf.yield %filled : tensor<?xf32>
+  }
+  %2 = "test.reify_bound"(%0) {type = "EQ", dim = 0} : (tensor<?xf32>) -> (index)
+  "test.some_use"(%2) : (index) -> ()
+  return
+}

mlir/test/lib/Dialect/Affine/CMakeLists.txt

add_mlir_library(MLIRAffineTransformsTestPasses
   Core
 
   LINK_LIBS PUBLIC
   MLIRAffineTransforms
   MLIRAffineUtils
   MLIRIR
   MLIRPass
   MLIRSupport
+  MLIRMemRefDialect
+  MLIRTensorDialect
   MLIRVectorUtils
   )

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

 //===- TestReifyValueBounds.cpp - Test value bounds reification -----------===//
 //
 // Part of the LLVM Project, 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 "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Affine/Transforms/Transforms.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Interfaces/ValueBoundsOpInterface.h"
 #include "mlir/Pass/Pass.h"
 
 #define PASS_NAME "test-affine-reify-value-bounds"
 
 using namespace mlir;
 using mlir::presburger::BoundType;
 
 namespace {
 
 /// This pass applies the permutation on the first maximal perfect nest.
 struct TestReifyValueBounds
     : public PassWrapper<TestReifyValueBounds, OperationPass<func::FuncOp>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestReifyValueBounds)
 
   StringRef getArgument() const final { return PASS_NAME; }
   StringRef getDescription() const final {
     return "Tests affine loop permutation utility";
   }
   TestReifyValueBounds() = default;
   TestReifyValueBounds(const TestReifyValueBounds &pass) : PassWrapper(pass){};
 
   void getDependentDialects(DialectRegistry &registry) const override {
-    registry.insert<AffineDialect>();
+    registry
+        .insert<AffineDialect, tensor::TensorDialect, memref::MemRefDialect>();
   }
 
   void runOnOperation() override;
 
 private:
   Option<bool> reifyToFuncArgs{
       *this, "reify-to-func-args",
       llvm::cl::desc("Reify in terms of function args"), llvm::cl::init(false)};
if (op->getName().getStringRef() == "test.reify_bound") {
                      : std::make_optional<int64_t>(
                            op->getAttrOfType<IntegerAttr>("dim").getInt());
 
       // 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) {
+      std::function<bool(Value, std::optional<int64_t>)> stopCondition =
+          [&](Value v, std::optional<int64_t> d) {
             // Reify in terms of SSA values that are different from `value`.
             return v != value;
           };
       if (reifyToFuncArgs) {
         // Reify in terms of function block arguments.
-        stopCondition = stopCondition = [](Value v) {
+        stopCondition = stopCondition = [](Value v, std::optional<int64_t> d) {
           auto bbArg = v.dyn_cast<BlockArgument>();
           if (!bbArg)
             return false;
           return isa<FunctionOpInterface>(
               bbArg.getParentBlock()->getParentOp());
         };
       }
 

utils/bazel/llvm-project-overlay/mlir/test/BUILD.bazel

deps = [
         "//mlir:Analysis",
         "//mlir:DialectUtils",
         "//mlir:FuncDialect",
         "//mlir:IR",
         "//mlir:MemRefDialect",
         "//mlir:Pass",
         "//mlir:SCFDialect",
         "//mlir:Support",
+        "//mlir:TensorDialect",
         "//mlir:Transforms",
         "//mlir:ValueBoundsOpInterface",
         "//mlir:VectorDialect",
         "//mlir:VectorUtils",
     ],
 )
 
 cc_library(