[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

86

doc about the arguments?

115

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

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

25

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

25

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/Interfaces/ValueBoundsOpInterface.h

private:
     ValueBoundsConstraintSet &cstr;
     Value value;
     std::optional<int64_t> dim = std::nullopt;
   };
 
 public:
   /// The stop condition when traversing the backward slice of a shaped value/
   /// index-type value. The traversal continues until the stop condition
   /// evaluates to "true" for a value.
-  using StopConditionFn = function_ref<bool(Value)>;
+  ///
+  /// The first parameter of the function is the shaped value/index-typed
+  /// value. The second parameter is the dimension in case of a shaped value.
+  using StopConditionFn =
+      function_ref<bool(Value, std::optional<int64_t> /*dim*/)>;
 
   /// 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.
   static LogicalResult computeBound(AffineMap &resultMap,
                                     ValueDimList &mapOperands,
                                     presburger::BoundType type, Value value,
                                     std::optional<int64_t> dim,
                                     StopConditionFn stopCondition);
 
+  /// Compute a bound in terms of the values/dimensions in `dependencies`. The
+  /// computed bound consists of only constant terms and dependent values (or
+  /// dimension sizes thereof).
+  static LogicalResult computeBound(AffineMap &resultMap,
+                                    ValueDimList &mapOperands,
+                                    presburger::BoundType type, Value value,
+                                    std::optional<int64_t> dim,
+                                    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

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

 using namespace mlir;
 
 FailureOr<OpFoldResult> mlir::reifyValueBound(OpBuilder &b, Location loc,
                                               presburger::BoundType type,
                                               Value value,
                                               std::optional<int64_t> dim) {
   // We are trying to reify a bound for `value`. Construct a stop condition that
   // evaluates to "true" for any SSA value expect for `value`. I.e., the bound
-  // will be computed in terms of any SSA values expect for `value`. The first
+  // will be computed in terms of any SSA values except for `value`. The first
   // such values are operands of the owner of `value`.
-  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

 //===- ValueBoundsOpInterfaceImpl.cpp - Impl. of ValueBoundsOpInterface ---===//
 //
 // 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/SCF/IR/ValueBoundsOpInterfaceImpl.h"
 
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Interfaces/ValueBoundsOpInterface.h"
 
 using namespace mlir;
+using presburger::BoundType;
 
 namespace mlir {
 namespace scf {
 namespace {
 
 struct ForOpInterface
     : public ValueBoundsOpInterface::ExternalModel<ForOpInterface, ForOp> {
+
+  /// Populate bounds of values/dimensions for iter_args/OpResults.
+  static void populateIterArgBounds(scf::ForOp forOp, Value value,
+                                    std::optional<int64_t> dim,
+                                    ValueBoundsConstraintSet &cstr) {
+    // `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 (dimension size)
+    // equals the corresponding block argument (dimension size).
+    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];
+
+    auto addEqBound = [&]() {
+      if (dim.has_value()) {
+        cstr.bound(value)[*dim] == cstr.getExpr(initArg, dim);
+      } else {
+        cstr.bound(value) == initArg;
+      }
+    };
+
+    if (yieldedValue == iterArg) {
+      addEqBound();
+      return;
+    }
+
+    // Compute EQ bound for yielded value.
+    AffineMap bound;
+    ValueDimList boundOperands;
+    LogicalResult status = ValueBoundsConstraintSet::computeBound(
+        bound, boundOperands, BoundType::EQ, yieldedValue, dim,
+        [&](Value v, std::optional<int64_t> d) {
+          // Stop when reaching a block argument of the loop body.
+          if (auto bbArg = v.dyn_cast<BlockArgument>())
+            return bbArg.getOwner()->getParentOp() == forOp;
+          // Stop when reaching a value that is defined outside of the loop. It
+          // is impossible to reach an iter_arg from there.
+          Operation *op = v.getDefiningOp();
+          return forOp.getLoopBody().findAncestorOpInRegion(*op) == nullptr;
+        });
+    if (failed(status))
+      return;
+    if (bound.getNumResults() != 1)
+      return;
+
+    // Check if computed bound equals the corresponding iter_arg.
+    Value singleValue = nullptr;
+    std::optional<int64_t> singleDim = std::nullopt;
+    if (auto dimExpr = bound.getResult(0).dyn_cast<AffineDimExpr>()) {
+      int64_t idx = dimExpr.getPosition();
+      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();
+      singleValue = boundOperands[idx].first;
+      singleDim = boundOperands[idx].second;
+    }
+    if (singleValue == iterArg && singleDim == dim)
+      addEqBound();
+  }
+
   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.bound(value) >= forOp.getLowerBound();
       cstr.bound(value) < forOp.getUpperBound();
+      return;
+    }
+
+    // Handle iter_args and OpResults.
+    populateIterArgBounds(forOp, value, std::nullopt, cstr);
   }
 
   void populateBoundsForShapedValueDim(Operation *op, Value value, int64_t dim,
                                        ValueBoundsConstraintSet &cstr) const {
-    // iter_arg / return value not supported.
-    return;
+    auto forOp = cast<ForOp>(op);
+    // Handle iter_args and OpResults.
+    populateIterArgBounds(forOp, value, dim, cstr);
   }
 };
 
 } // 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

if (dim != kIndexValue) {
       auto shapedType = cast<ShapedType>(value.getType());
       if (shapedType.hasRank() && !shapedType.isDynamicDim(dim)) {
         bound(value)[dim] == getExpr(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;
void ValueBoundsConstraintSet::projectOut(
   }
 }
 
 LogicalResult ValueBoundsConstraintSet::computeBound(
     AffineMap &resultMap, ValueDimList &mapOperands, presburger::BoundType type,
     Value value, std::optional<int64_t> dim, StopConditionFn stopCondition) {
 #ifndef NDEBUG
   assertValidValueDim(value, dim);
+  assert(!stopCondition(value, dim) &&
+         "stop condition should not be satisfied for starting point");
 #endif // NDEBUG
 
   Builder b(value.getContext());
   mapOperands.clear();
 
-  if (stopCondition(value)) {
+  if (stopCondition(value, dim)) {
     // Special case: If the stop condition is satisfied for the input
     // value/dimension, directly return it.
     mapOperands.push_back(std::make_pair(value, dim));
     resultMap = AffineMap::get(/*dimCount=*/1, /*symbolCount=*/0,
                                b.getAffineDimExpr(0));
     return success();
   }
 
   // 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(value, dim);
   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);
   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) {
+  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,
     StopConditionFn stopCondition) {
 #ifndef NDEBUG
   assertValidValueDim(value, dim);
 #endif // NDEBUG
 
   // Process the backward slice of `value` (i.e., reverse use-def chain) until
   // `stopCondition` is met.
   ValueBoundsConstraintSet cstr(value, dim);
   int64_t pos = cstr.getPos(value, dim);
   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
+}
+
+// -----
+
+func.func @scf_for_swapping_yield(%t1: tensor<?xf32>, %t2: tensor<?xf32>, %a: index, %b: index, %c: index) {
+  %cst = arith.constant 5.0 : f32
+  %r1, %r2 = scf.for %iv = %a to %b step %c iter_args(%arg1 = %t1, %arg2 = %t2) -> (tensor<?xf32>, tensor<?xf32>) {
+    %filled1 = linalg.fill ins(%cst : f32) outs(%arg1 : tensor<?xf32>) -> tensor<?xf32>
+    %filled2 = linalg.fill ins(%cst : f32) outs(%arg2 : tensor<?xf32>) -> tensor<?xf32>
+    scf.yield %filled2, %filled1 : tensor<?xf32>, tensor<?xf32>
+  }
+  // expected-error @below{{could not reify bound}}
+  %reify1 = "test.reify_bound"(%r1) {type = "EQ", dim = 0} : (tensor<?xf32>) -> (index)
+  "test.some_use"(%reify1) : (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 ValueBoundsOpInterface with affine dialect reification";
   }
   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" ||
                            op->getAttrOfType<IntegerAttr>("dim").getInt());
 
       // Check if a constant was requested.
       bool constant =
           op->getName().getStringRef() == "test.reify_constant_bound";
 
       // 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(