[mlir][Interfaces] Add ValueBoundsOpInterface and tensor dialect op impl

[springerm]

Summary

Ops can implement this interface to specify lower/upper bounds for their result values and block arguments. Bounds can be specified for:

• Index-type values
• Dimension sizes of shapes values

The bounds are added to a constraint set. Users can query this constraint set to compute bounds wrt. to a user-specified set of values. Only EQ bounds are supported at the moment.

A helper function is added to the Affine dialect to reify a computed bounds with an affine.apply op.

This revision also contains interface implementations for various tensor dialect ops, which illustrates how to implement this interface.

RFC: https://discourse.llvm.org/t/rfc-valueboundsopinterface-compute-bounds-of-index-values-dyn-dim-sizes/69174

Depends On: D146869

[springerm]

mlir/test/lib/Dialect/Linalg/TestLinalgTransforms.cpp

131

I have not found a good way yet to do this in the transform dialect. This change does not add a transform. It just adds an interface and its implementation. I think there are certain things that are easier (less code in the test file) to test via transform ops and other things are better to test via test passes.

I could add a transform.structured.reify_bounds op that takes a value handle and produces a value handle. This, combined with another new transform dialect op for "replaceAllUsesWith" could be a replacement for what I'm doing in TestLinalgTransforms.cpp. There would be a lot of boilerplate code though.

If this has to be a transform dialect op, ideally we would add an op that is also useful outside of tests. The main question is how are we going to use this functionality from the transform dialect. Do we have a use case for things such as linalg::computeUpperBoundForIndex as a transform op or this usually part of another transform op implementation? In the latter case, it may be better to keep this test in TestLinalgTransforms.cpp.

[dcaballe]

I took a first look and this looks really promising! Thank you so much for working on this so quickly.
Before we move forward with deeper reviews and comments, do you plan to send and RFC? I think this interface may benefit from one as it will impact multiple dialects. It would recommend so even if it's just an informative one.

Thanks!
Diego

mlir/include/mlir/Dialect/Linalg/Transforms/ValueBoundsOpInterface.h

24

Perhaps adding more context here about terms used in the public APIs would help (e.g., worklist, used value)

27

Don't have a mental model yet but this comment makes me think that this class holds a state (mapping) that may become invalid if the IR changes. Is that the case?

35

What about having to methods, one for values and one for dims? Clearer, less error-prone for the user and we wouldn't have to expose kIndexValue publicly.
Another option is using an std::optional. I'm not very excited about using a -1 but feel free to keep it if you think it simplifies the implementation/other APIs.

74

(Not sure I understand yet what a used value is)
It looks like this getExpr is returning an AffineExpr but also doing something else underneath. Make the name more explicit?

117

Would it make sense to build your own builder? If this is kind of an analysis, it may happen that the user doesn't have a builder. We should only need an MLIRContext.

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

1

typo

mlir/test/lib/Dialect/Linalg/TestLinalgTransforms.cpp

228

Doc about reifyToFuncArgs?
Doc about what this function does? I expected this to just print the bound but that's not the case.

[springerm]

In D145681#4186997, @dcaballe wrote:

Before we move forward with deeper reviews and comments, do you plan to send and RFC? I think this interface may benefit from one as it will impact multiple dialects. It would recommend so even if it's just an informative one.

Good idea: https://discourse.llvm.org/t/rfc-valueboundsopinterface-compute-bounds-of-index-values-dyn-dim-sizes/69174

mlir/include/mlir/Dialect/Linalg/Transforms/ValueBoundsOpInterface.h

27

Yes, that's correct. Also note that the only constructor of ValueBoundsConstraintSet is private and reifyBound does not return a ValueBoundsConstraintSet. I.e., there's no way for the user to get hold of a ValueBoundsConstraintSet object.

74

The argument name val may be confusing. This function simply returns an AffineConstExpr. Users could build one themselves without using this API, but then they would have to create a Builder. This is just a small convenience function.

117

This builder is created (from the MLIRContext object inside the ValueDim that is passed to the constructor) and owned by the ValueBoundsConstraintSet. It is just stored here so that we don't create a new builder every time an AffineExpr is created. Users cannot pass a builder to this class. (But they have to pass one to reifyBounds.)

mlir/test/lib/Dialect/Linalg/TestLinalgTransforms.cpp

228

Doc is in the Option definition in line 134: If set to true, values will be reified in terms of FuncOp bbArgs. (Otherwise, in terms of any values that are not owned by the same op; i.e., similar to reifyResultShapes.) Will add more docs to the reifyValueBounds function.

[mehdi_amini]

mlir/include/mlir/Dialect/Linalg/Transforms/ValueBoundsOpInterface.h

23

I don't know what "columns" refers to? I suspect some examples here on top of what Diego mentions below would help

27

I was about to ask the same thing: this is a map from SSA value to constraints, and it's not clear to me who holds this set and how to manage invalidation / recalculation?

44

mlir/include/mlir/Dialect/Linalg/Transforms/ValueBoundsOpInterface.td

14

Why is this inside Linalg? Is it because of "reifyBound" depending on tensor/memref dialects? (see my comment there)

This is an major red flag to me right now.

30

It's not clear to me what the "value" is here, is this a result of the operation? The doc does not seem to say.

In this case can we pass an index instead? Or an OpResult?

Edit: looking at the implementation, it can also be the parent op in case this is a block argument! I suspect the current implementation does not seem suitable for a CFG?

mlir/include/mlir/InitAllDialects.h

136

There were supposed to be a mechanism to make this all safe that was discussed when you started using this pattern initially, was it implemented?

mlir/lib/Dialect/Linalg/Transforms/ValueBoundsOpInterface.cpp

303

Can we replace this with an interface method on ShapedTypeInterface and avoid hard-depending on specific dialect implementations?

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

2

Please don't use -allow-unregistered-dialect anywhere, we should never use this in the codebase.

[springerm]

mlir/include/mlir/Dialect/Linalg/Transforms/ValueBoundsOpInterface.td

14

No particular reason. It should not be in LinalgTransforms, I just had to put it somewhere for the moment.

My preferred solution would be:

• Use IntegerRelation instead of FlatAffineValueConstraints. (IntegerRelation/IntegerPolyhedron needs a few functions that are currently defined on FlatAffineValueConstraints, but I think this is doable.) ValueBoundsOpInterface and ValueBoundsConstraintSet can then be in mlir/Interfaces.
• ValueBoundsConstraintSet::reifyBound becomes a standalone static function in AffineTransforms (or AffineDialect if we can find a way to create memref.dim/tensor.dim without depending on TensorDialect/MemRefDialect).
• We could add a second reifyBound function to ArithTransforms (that will create std.addi etc. instead of affine.apply).

30

In most cases it would be an OpResult. There are a few cases where we would want to support a BlockArgument. E.g., it could be the IV of an scf.for. I think in case of a BlockArgument, this should be restricted to "entry" block arguments of regions.

(I think the current worklist approach is too simple to handle cyclic CFGs etc.; it could be extended in the future if there are use cases. Same goes for iter_args of scf.for etc.: D145804 adds an interface impl for scf.for, but iter_args are not supported.)

I will update the docs accordingly.

[dcaballe]

mlir/include/mlir/Dialect/Linalg/Transforms/ValueBoundsOpInterface.h

27

Same question as Mehdi.

This looks like an analysis to me. I haven't seen analysis passes implemented in MLIR but we have the infra: https://mlir.llvm.org/docs/PassManagement/#analysis-management. I wonder if this would be a good fit for an analysis.

117

(But they have to pass one to reifyBounds.)

Exactly. I had to create a dummy builder just to invoke reifyBounds.

[mehdi_amini]

mlir/include/mlir/InitAllDialects.h

136

Found the revision: https://reviews.llvm.org/D120368

I would really like to see this getting through before we add more "external interface" in an unsafe fashion. Can you help with this?

[springerm]

mlir/include/mlir/Dialect/Linalg/Transforms/ValueBoundsOpInterface.h

35

std::optional did not work the DenseMap, but we can at least use it for the public API.

[springerm]

mlir/include/mlir/Dialect/Linalg/Transforms/ValueBoundsOpInterface.h

117

That is because reifyBounds can be used to reify non-constant bounds. I.e., the bound may be the result of an affine.apply op that computes the bound. You're using it to compute constant bounds (in which case no builder would be needed), but reifyBounds is more general.

I'm going to add a new option/overload for computing an explicitly constant bound (without a builder) in a subsequent revision.

mlir/lib/Dialect/Linalg/Transforms/ValueBoundsOpInterface.cpp

303

D146052

[springerm]

address some comments and move to mlir/Interfaces

[springerm]

update

[springerm]

rebase

[springerm]

Remove loc

[springerm]

mlir/include/mlir/InitAllDialects.h

136

D146414

[springerm]

rebase

[dcaballe]

LGTM. All my comments have been resolved. Please, wait for Mehdi to provide feedback on the remaining open questions/secondary PRs.

[mehdi_amini]

I'm curious about @jpienaar 's thoughts on the overall direction here

mlir/include/mlir/InitAllDialects.h

136

I am not sure what you have in mind for D146414, I meant helping reviewing D120368 to get it to land.

mlir/lib/Analysis/FlatLinearValueConstraints.cpp

460 ↗

(On Diff #506945)

Can you commit this separately? It's not clear to me that it relates to the current patch

[springerm]

rebase

[jpienaar]

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

52 ↗

(On Diff #508292)

So value and dim may be provided as long as value is not index type? E.g., if value is an i32, then you can have both?

59 ↗

(On Diff #508292)

It is difficult to know what this function will do and stopCondition represents - it feels like base case of recursion, but its unclear here. Could you perhaps expand this to provide more information? How many stop conditions are actually here?

64 ↗

(On Diff #508292)

And ValueBoundsConstraintSet is not visible here so that one could use the stop condition type defined there?

mlir/include/mlir/Interfaces/ValueBoundsOpInterface.h

36 ↗

(On Diff #508292)

That's just being conservative right? Could this be retained in more cases?

44 ↗

(On Diff #508292)

OOC are all BoundType's allowed?

61 ↗

(On Diff #508292)

I almost wonder if spelling out boundtype in name isn't more intuitive, e.g.,

addUpperBound, addLowerBound, addEquality (last one sounds less good than other two).

Ha, makes me wonder about operators boundset.bound(value) <= dim / boundset.bound(value)[dim] <= expr - which does also translate closer to

addBound(value, UB, expr) - so it reads add bound for value of upper bound with expr.

Seeing presburger in the API just feels weird to me.

79 ↗

(On Diff #508292)

These aren't bounds right? This feels natural with equality bounds/solution of constraint set.

91 ↗

(On Diff #508292)

Can we choose any other sentinel (-56) or differentiate with type (optional)? (-1 has been used in sentinel in too many places ...)

101 ↗

(On Diff #508292)

Mehdi commented on this as well, but this is first mention of column, might be worthwhile to introduce here.

147 ↗

(On Diff #508292)

Could this be along with destination style ops or in another header?

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

23 ↗

(On Diff #508292)

Could you expand here? Its not quite intuitive what the effect is (well and presents a nicely documented usage of the interface).

53 ↗

(On Diff #508292)

New form is

cast<Foo>(...) / isa<Foo>(...)

(https://mlir.llvm.org/deprecation/)

57 ↗

(On Diff #508292)

(your other change will improve this, I also toyed around with having a shaped.dim op that canonicalizes on input type without dependency on dialect https://gist.github.com/jpienaar/82c79ae73b2f46be4c9f365334bf8dc3)

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

19 ↗

(On Diff #508292)

I'd really prefer if we did this in a more declarative manner so that this could be generated and we could even have a dynamic form of this (similar to how torch-mlir does shape functions in spirit). But more work is needed for that and we aren't there yet, and that can't be a blocker here.

73 ↗

(On Diff #508292)

This feels like exactly what we'd represent in the shape (and type inference) function.

119 ↗

(On Diff #508292)

Why declare this as an external interface inside the dialect directory? Are alternative implementations expected? This going to be optionally added?

mlir/lib/Interfaces/ValueBoundsOpInterface.cpp

64 ↗

(On Diff #508292)

Could we capture the error in op and surface it when queried or some such? Or is it assumed that verification would have caught all instances where bounds could fail?

149 ↗

(On Diff #508292)

https://llvm.org/docs/CodingStandards.html#don-t-evaluate-end-every-time-through-a-loop

[springerm]

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

52 ↗

(On Diff #508292)

i32 is not supported (added with D146870). Only index-type values and shaped values are supported.

I added an assertion in ValueBoundsConstraintSet::computeBound.

59 ↗

(On Diff #508292)

Added an example. Some of the next revisions also add overloads, so that users can use this interface without specifying a stopCondition.

• D146296 makes the stopCondition optional for constant bounds.
• D146306 adds an overload for a given list of dependencies.
• D143910 adds an overload for a given list of independencies.

mlir/include/mlir/Interfaces/ValueBoundsOpInterface.h

36 ↗

(On Diff #508292)

Correct. Adding new IR is OK. Expanded the comment.

44 ↗

(On Diff #508292)

Only EQ at the moment. There is an assertion in the implementation that LB/UB are not yet implemented. The other bound types are added in D145787 to keep this revision a bit smaller.

61 ↗

(On Diff #508292)

This is a nice API. And it allows us to specify inclusive/exclusive bounds via </<= etc.

I updated just this revision for now, will update the dependent ones later.

79 ↗

(On Diff #508292)

The getExpr function can be used to provide an expression that can be passed to addBound. (The overload that take an AffineExpr.)

91 ↗

(On Diff #508292)

I wanted to use std::optional but for some reason, an std::optional cannot be used as the key for a DenseMap:

DenseMap.h:469:22: error: no member named 'getEmptyKey' in 'llvm::DenseMapInfo<std::optional<long>>'

147 ↗

(On Diff #508292)

This must be in some place where both DestinationStyleOpInterface and ValueBoundsOpInterface are available. (Kind of like cross-dialect dependencies, but this one is cross-interfaces.)

ValueBoundsOpInterface seems a bit better than putting it in DestinationStyleOpInterface.h because when DestinationStyleOpInterface is a more "fundamental" interface.

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

19 ↗

(On Diff #508292)

What do you mean by "more declarative"?

119 ↗

(On Diff #508292)

It could be implemented directly on the op. Alternative interface implementations are not possible in MLIR (not even with external models).

The reason why I put it here is separation of concerns. So that TensorOps.cpp/TensorOps.td is not getting any longer. This keeps TensorOps.td clean of any "optional/non-fundamental" interfaces and easy to understand. I think certain interfaces are very important and should be implemented on the op directly (e.g., OffsetSizeAndStrideOpInterface), other just clutter the file (like this one).

mlir/lib/Interfaces/ValueBoundsOpInterface.cpp

64 ↗

(On Diff #508292)

This is just a sanity check. This should always return success. If it fails, it indicates that there is a bug in FlatLinearConstraints or the ValueBoundsOpInterface implementation.

[springerm]

address most comments

[springerm]

update

[springerm]

mlir/include/mlir/Interfaces/ValueBoundsOpInterface.h

91 ↗

(On Diff #508292)

I updated a few more places in the internal API. kIndexValue is now just used when interacting with the DenseMap.

[jpienaar]

Nice, I like the new API too.

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

52 ↗

(On Diff #508292)

I meant something else: is this optional simulating a sum type? E.g., is only one of value or dim allowed to be set and an optional is being used? So a ValueOrDim bound is actually being passed in using more raw types (well a std::variant<Value, int64_t> is probably not much more higher level than std::optional, so raw perhaps not best word)

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

19 ↗

(On Diff #508292)

Something that could allow us to capture the information in one way but use both places rather than implement twice. I feel like if we had it such that you could run an abstract interpretation over the shape inference functions, you could have gotten some of these for free (and vice versa). [I'll probably prod at this in a month or so, so no blocker here, just me wishing that was done already :)]

mlir/lib/Interfaces/ValueBoundsOpInterface.cpp

49 ↗

(On Diff #509962)

Following new style preference/deprecation, dyn_cast<ShapedType>(value.getType()) here and above.

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

29 ↗

(On Diff #509962)

Update description ?

[springerm]

Thanks for the thorough review!

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

52 ↗

(On Diff #508292)

std::variant is like a union right? This is different here. This function can be called in two ways:

• value has index type and dim = nullopt
• value has a shaped type and dim != nullopt

So value is always needed and whether dim should be specified or not depends on value.getType().

mlir/include/mlir/Dialect/Linalg/CMakeLists.txt

 add_subdirectory(IR)
 add_subdirectory(TransformOps)
+add_subdirectory(Transforms)
 
 set(LLVM_TARGET_DEFINITIONS Passes.td)
 mlir_tablegen(Passes.h.inc -gen-pass-decls -name Linalg)
 mlir_tablegen(Passes.capi.h.inc -gen-pass-capi-header --prefix Linalg)
 mlir_tablegen(Passes.capi.cpp.inc -gen-pass-capi-impl --prefix Linalg)
 add_public_tablegen_target(MLIRLinalgPassIncGen)
 
 add_mlir_doc(Passes LinalgPasses ./ -gen-pass-doc)

mlir/include/mlir/Dialect/Linalg/Transforms/CMakeLists.txt

+add_mlir_interface(ValueBoundsOpInterface)

mlir/include/mlir/Dialect/Linalg/Transforms/ValueBoundsOpInterface.h

+//===- ValueBoundsOpInterface.h - Value Bounds ------------------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_LINALG_TRANSFORMS_VALUEBOUNDSOPINTERFACE_H_
+#define MLIR_DIALECT_LINALG_TRANSFORMS_VALUEBOUNDSOPINTERFACE_H_
+
+#include "mlir/Dialect/Affine/Analysis/AffineStructures.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"
+
+namespace mlir {
+namespace linalg {
+
+/// A helper class to be used with `ValueBoundsOpInterface`. This class stores a
+/// constraint system and mapping of columns to values/shape dimensions.
+///
+/// Note: This class maintains its own mapping to SSA values; no SSA values are
+/// mapped in the underlying `FlatAffineValueConstraints`. This is because not
+/// only SSA values but also shape dimensions of SSA values must be mapped.
+class ValueBoundsConstraintSet {
+public:
+  /// Reify a bound for the given index-typed value or shape dimension size in
+  /// terms of the owning op's operands.
+  static FailureOr<OpFoldResult>
+  reifyBound(OpBuilder &b, Location loc,
+             presburger::IntegerPolyhedron::BoundType type, Value value,
+             int64_t dim = kIndexValue);
+
+  /// Reify a bound for the given index-typed value or shape dimension size in
+  /// terms of SSA values for which `stopCondition` is met.
+  static FailureOr<OpFoldResult>
+  reifyBound(OpBuilder &b, Location loc,
+             presburger::IntegerPolyhedron::BoundType type, Value value,
+             int64_t dim, function_ref<bool(Value)> stopCondition);
+
+  /// Dimension indentifier to indicate a value is index-typed.
+  static const int64_t kIndexValue = -1;
+
+  /// Bound the given index-typed value by the given expression.
+  void addBound(presburger::IntegerPolyhedron::BoundType type, Value value,
+                AffineExpr expr);
+
+  /// Bound the the given shaped value dimension by the given expression.
+  void addBound(presburger::IntegerPolyhedron::BoundType type, Value value,
+                int64_t dim, AffineExpr expr);
+
+  /// Bound the given index-typed value by a constant or SSA value.
+  void addBound(presburger::IntegerPolyhedron::BoundType type, Value value,
+                OpFoldResult ofr);
+
+  /// Bound the the given shaped value dimension by a constant or SSA value.
+  void addBound(presburger::IntegerPolyhedron::BoundType type, Value value,
+                int64_t dim, OpFoldResult ofr);
+
+  /// Return an expression that represents the given index-typed value or shaped
+  /// value dimension. If this value/dimension was not used so far, it is added
+  /// to the worklist.
+  AffineExpr getExpr(Value value, int64_t dim = kIndexValue);
+
+  /// Return an expression that represents a constant or index-typed SSA value.
+  /// In case of a value, if this value was not used so far, it is added to the
+  /// worklist.
+  AffineExpr getExpr(OpFoldResult ofr);
+
+  /// Return an expression that represents a constant.
+  AffineExpr getExpr(int64_t val);
+
+private:
+  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);
+
+  /// Bound the given column in the underlying constraint set by the given
+  /// expression.
+  void addBound(presburger::IntegerPolyhedron::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;
+
+  /// Insert a value/dimension into the constraint set. If `isSymbol` is set to
+  /// "false", a dimension is added.
+  ///
+  /// Note: There are certain affine restrictions wrt. dimensions. E.g., they
+  /// cannot be multiplied. Furthermore, bounds can only be queried for
+  /// dimensions but not for symbols.
+  int64_t insert(ValueDim valueDim, bool isSymbol = true);
+
+  /// Project out the given column in the constraint set.
+  void projectOut(int64_t pos);
+
+  /// 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;
+
+  /// Constraint system of equalities and inequalities.
+  FlatAffineValueConstraints cstr;
+
+  /// Builder for constructing affine expressions.
+  Builder builder;
+};
+} // namespace linalg
+} // namespace mlir
+
+#include "mlir/Dialect/Linalg/Transforms/ValueBoundsOpInterface.h.inc"
+
+namespace mlir {
+namespace linalg {
+
+/// Default implementation for destination style ops: Tied OpResults and
+/// OpOperands have the same type.
+template <typename ConcreteOp>
+struct DstValueBoundsOpInterfaceExternalModel
+    : public ValueBoundsOpInterface::ExternalModel<
+          DstValueBoundsOpInterfaceExternalModel<ConcreteOp>, ConcreteOp> {
+  void populateBoundsForShapedValueDim(Operation *op, Value value, int64_t dim,
+                                       ValueBoundsConstraintSet &cstr) const {
+    auto dstOp = cast<DestinationStyleOpInterface>(op);
+    assert(value.getDefiningOp() == dstOp);
+
+    Value tiedOperand = dstOp.getTiedOpOperand(value.cast<OpResult>())->get();
+    cstr.addBound(presburger::IntegerPolyhedron::BoundType::EQ, value, dim,
+                  cstr.getExpr(tiedOperand, dim));
+  }
+};
+
+} // namespace linalg
+} // namespace mlir
+
+#endif // MLIR_DIALECT_LINALG_TRANSFORMS_VALUEBOUNDSOPINTERFACE_H_

mlir/include/mlir/Dialect/Linalg/Transforms/ValueBoundsOpInterface.td

+//===-- ValueBoundsOpInterface.td - Value Bounds -----------*- tablegen -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef VALUEBOUNDSOPINTERFACE
+#define VALUEBOUNDSOPINTERFACE
+
+include "mlir/IR/OpBase.td"
+
+def ValueBoundsOpInterface : OpInterface<"ValueBoundsOpInterface"> {
+  let description = [{
+    This interface allows operations with index-typed and/or shaped value-typed
+    results/block arguments to specify range bounds. These bounds are stored in
+    a constraint set. The constraint set can then be queried to compute bounds
+    in terms of other values that are stored in the constraint set.
+  }];
+  let cppNamespace = "::mlir::linalg";
+  let methods = [
+    InterfaceMethod<
+      /*desc=*/[{
+        Populate the constraint set with bounds for the given index-typed
+        value.
+      }],
+      /*retType=*/"void",
+      /*methodName=*/"populateBoundsForIndexValue",
+      /*args=*/(ins "::mlir::Value":$value,
+                    "::mlir::linalg::ValueBoundsConstraintSet &":$cstr),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        llvm_unreachable("not implemented");
+      }]
+    >,
+    InterfaceMethod<
+      /*desc=*/[{
+        Populate the constraint set with bounds for the size of the specified
+        dimension of the given shaped value.
+      }],
+      /*retType=*/"void",
+      /*methodName=*/"populateBoundsForShapedValueDim",
+      /*args=*/(ins "::mlir::Value":$value,
+                    "int64_t":$dim,
+                    "::mlir::linalg::ValueBoundsConstraintSet &":$cstr),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        llvm_unreachable("not implemented");
+      }]
+    >,
+  ];
+}
+
+#endif  // VALUEBOUNDSOPINTERFACE

mlir/include/mlir/Dialect/Linalg/Transforms/ValueBoundsOpInterfaceImpl.h

+//===- ValueBoundsOpInterfaceImpl.h - 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_LINALG_TRANSFORMS_VALUEBOUNDSOPINTERFACEIMPL_H
+#define MLIR_DIALECT_LINALG_TRANSFORMS_VALUEBOUNDSOPINTERFACEIMPL_H
+
+namespace mlir {
+class DialectRegistry;
+
+namespace linalg {
+void registerValueBoundsOpInterfaceExternalModels(DialectRegistry &registry);
+} // namespace linalg
+} // namespace mlir
+
+#endif // MLIR_DIALECT_LINALG_TRANSFORMS_VALUEBOUNDSOPINTERFACEIMPL_H

mlir/include/mlir/InitAllDialects.h

 #include "mlir/Dialect/Index/IR/IndexDialect.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/LLVMIR/NVVMDialect.h"
 #include "mlir/Dialect/LLVMIR/ROCDLDialect.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.h"
 #include "mlir/Dialect/Linalg/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/Linalg/Transforms/TilingInterfaceImpl.h"
+#include "mlir/Dialect/Linalg/Transforms/ValueBoundsOpInterfaceImpl.h"
 #include "mlir/Dialect/MLProgram/IR/MLProgram.h"
 #include "mlir/Dialect/Math/IR/Math.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/MemRef/TransformOps/MemRefTransformOps.h"
 #include "mlir/Dialect/MemRef/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/MemRef/Transforms/RuntimeOpVerification.h"
 #include "mlir/Dialect/NVGPU/IR/NVGPUDialect.h"
 #include "mlir/Dialect/OpenACC/OpenACC.h"
inline void registerAllDialects(DialectRegistry &registry) {
   vector::registerTransformDialectExtension(registry);
 
   // Register all external models.
   arith::registerBufferizableOpInterfaceExternalModels(registry);
   bufferization::func_ext::registerBufferizableOpInterfaceExternalModels(
       registry);
   linalg::registerBufferizableOpInterfaceExternalModels(registry);
   linalg::registerTilingInterfaceExternalModels(registry);
+  linalg::registerValueBoundsOpInterfaceExternalModels(registry);
   memref::registerBufferizableOpInterfaceExternalModels(registry);
   memref::registerRuntimeVerifiableOpInterfaceExternalModels(registry);
   scf::registerBufferizableOpInterfaceExternalModels(registry);
   shape::registerBufferizableOpInterfaceExternalModels(registry);
   sparse_tensor::registerBufferizableOpInterfaceExternalModels(registry);
   tensor::registerBufferizableOpInterfaceExternalModels(registry);
   tensor::registerInferTypeOpInterfaceExternalModels(registry);
   tensor::registerTilingInterfaceExternalModels(registry);

mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp

 /// variables (starting at 'offset') as affine maps of the remaining
 /// variables (dimensional and symbolic variables). Local variables are
 /// themselves explicitly computed as affine functions of other variables in
 /// this process if needed.
 void FlatAffineValueConstraints::getSliceBounds(
     unsigned offset, unsigned num, MLIRContext *context,
     SmallVectorImpl<AffineMap> *lbMaps, SmallVectorImpl<AffineMap> *ubMaps,
     bool getClosedUB) {
-  assert(num < getNumDimVars() && "invalid range");
+  assert(offset + num <= getNumDimVars() && "invalid range");
 
   // Basic simplification.
   normalizeConstraintsByGCD();
 
   LLVM_DEBUG(llvm::dbgs() << "getSliceBounds for first " << num
                           << " variables\n");
   LLVM_DEBUG(dump());
 

mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt

add_mlir_dialect_library(MLIRLinalgTransforms
   Promotion.cpp
   Split.cpp
   SplitReduction.cpp
   SubsetHoisting.cpp
   SwapExtractSliceWithFillPatterns.cpp
   Tiling.cpp
   TilingInterfaceImpl.cpp
   Transforms.cpp
+  ValueBoundsOpInterface.cpp
+  ValueBoundsOpInterfaceImpl.cpp
   Vectorization.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Linalg
 
   DEPENDS
   MLIRLinalgPassIncGen
 
   LINK_LIBS PUBLIC
+  MLIRAffineAnalysis
   MLIRAffineDialect
   MLIRAffineUtils
   MLIRAnalysis
   MLIRArithDialect
   MLIRArithTransforms
   MLIRArithUtils
   MLIRBufferizationDialect
   MLIRBufferizationTransforms

mlir/lib/Dialect/Linalg/Transforms/ValueBoundsOpInterface.cpp

+//===- ValueBoundsOpInterface.cpp - Value Bounds  -------------------------===//
+//
+// 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/Linalg/Transforms/ValueBoundsOpInterface.h"
+
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/Utils/StaticValueUtils.h"
+#include "mlir/IR/BuiltinTypes.h"
+
+using namespace mlir;
+using namespace mlir::linalg;
+using presburger::IntegerPolyhedron;
+
+namespace mlir {
+#include "mlir/Dialect/Linalg/Transforms/ValueBoundsOpInterface.cpp.inc"
+} // namespace mlir
+
+ValueBoundsConstraintSet::ValueBoundsConstraintSet(ValueDim valueDim)
+    : builder(valueDim.first.getContext()) {
+  insert(valueDim, /*isSymbol=*/false);
+}
+
+#ifndef NDEBUG
+static void assertValidValueDim(Value value, int64_t dim) {
+  if (value.getType().isIndex()) {
+    assert(dim == ValueBoundsConstraintSet::kIndexValue && "invalid dim value");
+  } else if (auto shapedType = value.getType().dyn_cast<ShapedType>()) {
+    assert(dim >= 0 && "invalid dim value");
+    if (shapedType.hasRank())
+      assert(dim < shapedType.getRank() && "invalid dim value");
+  } else {
+    llvm_unreachable("unsupported type");
+  }
+}
+#endif // NDEBUG
+
+void ValueBoundsConstraintSet::addBound(
+    presburger::IntegerPolyhedron::BoundType type, int64_t pos,
+    AffineExpr expr) {
+  LogicalResult status = cstr.addBound(
+      type, pos,
+      AffineMap::get(cstr.getNumDimVars(), cstr.getNumSymbolVars(), expr));
+  (void)status;
+  assert(succeeded(status) && "failed to add bound to constraint system");
+}
+
+void ValueBoundsConstraintSet::addBound(
+    presburger::IntegerPolyhedron::BoundType type, Value value,
+    AffineExpr expr) {
+  assert(value.getType().isIndex() && "expected index type");
+  addBound(type, getPos(value), expr);
+}
+
+void ValueBoundsConstraintSet::addBound(
+    presburger::IntegerPolyhedron::BoundType type, Value value, int64_t dim,
+    AffineExpr expr) {
+#ifndef NDEBUG
+  assertValidValueDim(value, dim);
+#endif // NDEBUG
+  addBound(type, getPos(value, dim), expr);
+}
+
+void ValueBoundsConstraintSet::addBound(
+    presburger::IntegerPolyhedron::BoundType type, Value value,
+    OpFoldResult ofr) {
+  assert(value.getType().isIndex() && "expected index type");
+  addBound(type, getPos(value), getExpr(ofr));
+}
+
+void ValueBoundsConstraintSet::addBound(
+    presburger::IntegerPolyhedron::BoundType type, Value value, int64_t dim,
+    OpFoldResult ofr) {
+#ifndef NDEBUG
+  assertValidValueDim(value, dim);
+#endif // NDEBUG
+  addBound(type, getPos(value, dim), getExpr(ofr));
+}
+
+AffineExpr ValueBoundsConstraintSet::getExpr(Value value, int64_t dim) {
+#ifndef NDEBUG
+  assertValidValueDim(value, dim);
+#endif // NDEBUG
+
+  auto shapedType = value.getType().dyn_cast<ShapedType>();
+  if (shapedType) {
+    // Static dimension: return constant directly.
+    if (shapedType.hasRank() && !shapedType.isDynamicDim(dim))
+      return builder.getAffineConstantExpr(shapedType.getDimSize(dim));
+  } else {
+    // Constant index value: return directly.
+    if (auto constInt = getConstantIntValue(value))
+      return builder.getAffineConstantExpr(*constInt);
+  }
+
+  // Dynamic value: add to constraint set.
+  ValueDim valueDim = std::make_pair(value, dim);
+  if (valueDimToPosition.find(valueDim) == valueDimToPosition.end())
+    (void)insert(valueDim);
+  int64_t pos = getPos(value, dim);
+  return pos < cstr.getNumDimVars()
+             ? builder.getAffineDimExpr(pos)
+             : builder.getAffineSymbolExpr(pos - cstr.getNumDimVars());
+}
+
+AffineExpr ValueBoundsConstraintSet::getExpr(OpFoldResult ofr) {
+  if (Value value = ofr.dyn_cast<Value>())
+    return getExpr(value, /*dim=*/kIndexValue);
+  auto constInt = getConstantIntValue(ofr);
+  assert(constInt.has_value() && "expected Integer constant");
+  return builder.getAffineConstantExpr(*constInt);
+}
+
+AffineExpr ValueBoundsConstraintSet::getExpr(int64_t val) {
+  return builder.getAffineConstantExpr(val);
+}
+
+int64_t ValueBoundsConstraintSet::insert(ValueDim valueDim, bool isSymbol) {
+  assert((valueDimToPosition.find(valueDim) == valueDimToPosition.end()) &&
+         "already mapped");
+  int64_t pos = isSymbol ? cstr.appendSymbolVar() : cstr.appendDimVar();
+  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);
+  return pos;
+}
+
+int64_t ValueBoundsConstraintSet::getPos(Value value, int64_t dim) const {
+#ifndef NDEBUG
+  assertValidValueDim(value, 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();
+    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(IntegerPolyhedron::BoundType::EQ, value, dim,
+                 builder.getAffineConstantExpr(shapedType.getDimSize(dim)));
+        continue;
+      }
+    }
+
+    // Do not process any further if the stop condition is met.
+    if (stopCondition(value))
+      continue;
+
+    // Query `ValueBoundsOpInterface` for constraints. New items may be added to
+    // the worklist.
+    auto valueBoundsOp =
+        dyn_cast<ValueBoundsOpInterface>(getOwnerOfValue(value));
+    if (!valueBoundsOp)
+      continue;
+    if (dim == kIndexValue) {
+      valueBoundsOp.populateBoundsForIndexValue(value, *this);
+    } else {
+      valueBoundsOp.populateBoundsForShapedValueDim(value, dim, *this);
+    }
+  }
+}
+
+void ValueBoundsConstraintSet::projectOut(int64_t pos) {
+  cstr.projectOut(pos);
+  bool erased = valueDimToPosition.erase(positionToValueDim[pos]);
+  assert(erased && "inconsistent reverse mapping");
+  positionToValueDim.erase(positionToValueDim.begin() + pos);
+  // Update reverse mapping.
+  for (int64_t i = pos; i < positionToValueDim.size(); ++i)
+    valueDimToPosition[positionToValueDim[i]] = i;
+}
+
+FailureOr<OpFoldResult> ValueBoundsConstraintSet::reifyBound(
+    OpBuilder &b, Location loc, presburger::IntegerPolyhedron::BoundType type,
+    Value value, int64_t dim) {
+  auto stopCondition = [&](Value v) {
+    // Reify in terms of SSA values that are different from `value`.
+    return v != value;
+  };
+  return ValueBoundsConstraintSet::reifyBound(b, loc, type, value, dim,
+                                              stopCondition);
+}
+
+FailureOr<OpFoldResult> ValueBoundsConstraintSet::reifyBound(
+    OpBuilder &b, Location loc, presburger::IntegerPolyhedron::BoundType type,
+    Value value, int64_t dim, function_ref<bool(Value)> stopCondition) {
+#ifndef NDEBUG
+  assertValidValueDim(value, dim);
+#endif // NDEBUG
+
+  // Only EQ bounds are supported at the moment.
+  assert(type == presburger::IntegerPolyhedron::BoundType::EQ &&
+         "unsupported bound type");
+
+  // Process the backward slice of `value` (i.e., reverse use-def chain) until
+  // `stopCondition` is met.
+  ValueBoundsConstraintSet cstr(std::make_pair(value, dim));
+  int64_t pos = cstr.getPos(value, dim);
+  cstr.processWorklist(stopCondition);
+
+  // Project out all positions (apart from `pos`) that do not match the stop
+  // condition.
+  int64_t nextPos = 0;
+  while (nextPos < cstr.positionToValueDim.size()) {
+    if (nextPos == pos) {
+      ++nextPos;
+      continue;
+    }
+
+    if (!stopCondition(cstr.positionToValueDim[nextPos].first)) {
+      cstr.projectOut(nextPos);
+      // The column was projected out so another column is now at that position.
+      // Do not increase the counter.
+    } else {
+      ++nextPos;
+    }
+  }
+
+  // Compute lower and upper bounds for `value`.
+  SmallVector<AffineMap> lb(1), ub(1);
+  cstr.cstr.getSliceBounds(pos, 1, b.getContext(), &lb, &ub,
+                           /*getClosedUB=*/true);
+  // Note: There are TODOs in the implementation of `getSliceBounds`. In such a
+  // case, no lower/upper bound can be computed at the moment.
+  if (lb.empty() || !lb[0] || ub.empty() || !ub[0] ||
+      lb[0].getNumResults() != 1 || ub[0].getNumResults() != 1)
+    return failure();
+
+  // Look for same lower and upper bound: EQ bound.
+  if (ub[0] != lb[0])
+    return failure();
+
+  // Gather all SSA values that are used in the computed bound.
+  SmallVector<Value> operands;
+  assert(cstr.cstr.getNumDimAndSymbolVars() == cstr.positionToValueDim.size() &&
+         "inconsistent mapping state");
+  for (int64_t i = 0; i < cstr.cstr.getNumDimAndSymbolVars(); ++i) {
+    // Skip `value`.
+    if (i == pos)
+      continue;
+    // Check if the position `i` is used in the generated bound. If so, it must
+    // be included in the generated affine.apply op.
+    bool used = false;
+    if (i < cstr.cstr.getNumDimVars()) {
+      if (lb[0].isFunctionOfDim(i))
+        used = true;
+    } else {
+      if (lb[0].isFunctionOfSymbol(i - cstr.cstr.getNumDimVars()))
+        used = true;
+    }
+
+    if (!used) {
+      // Not used: Put an empty Value (will canonicalize away).
+      operands.push_back(Value());
+      continue;
+    }
+
+    ValueBoundsConstraintSet::ValueDim valueDim = cstr.positionToValueDim[i];
+    Value value = valueDim.first;
+    int64_t dim = valueDim.second;
+    if (dim == ValueBoundsConstraintSet::kIndexValue) {
+      // An index-type value is used: can be used directly in the affine.apply
+      // op.
+      assert(value.getType().isIndex() && "expected index type");
+      operands.push_back(value);
+      continue;
+    }
+
+    assert(value.getType().cast<ShapedType>().isDynamicDim(dim) &&
+           "expected dynamic dim");
+    if (value.getType().isa<RankedTensorType>()) {
+      // A tensor dimension is used: generate a tensor.dim.
+      operands.push_back(b.create<tensor::DimOp>(loc, value, dim));
+    } else if (value.getType().isa<MemRefType>()) {
+      // A memref dimension is used: generate a memref.dim.
+      operands.push_back(b.create<memref::DimOp>(loc, value, dim));
+    } else {
+      llvm_unreachable("cannot generate DimOp for unsupported shaped type");
+    }
+  }
+
+  mlir::canonicalizeMapAndOperands(&lb[0], &operands);
+  // Check for special cases where no affine.apply op is needed.
+  if (lb[0].isSingleConstant()) {
+    // Bound is a constant: return an IntegerAttr.
+    return static_cast<OpFoldResult>(
+        b.getIndexAttr(lb[0].getSingleConstantResult()));
+  }
+  // No affine.apply op is needed if the bound is a single SSA value.
+  if (auto expr = lb[0].getResult(0).dyn_cast<AffineDimExpr>())
+    return static_cast<OpFoldResult>(operands[expr.getPosition()]);
+  if (auto expr = lb[0].getResult(0).dyn_cast<AffineSymbolExpr>())
+    return static_cast<OpFoldResult>(
+        operands[expr.getPosition() + cstr.cstr.getNumDimVars() - 1]);
+  // General case: build affine.apply op.
+  return static_cast<OpFoldResult>(
+      b.create<AffineApplyOp>(loc, lb[0], operands).getResult());
+}

mlir/lib/Dialect/Linalg/Transforms/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/Linalg/Transforms/ValueBoundsOpInterfaceImpl.h"
+
+#include "mlir/Dialect/Affine/Analysis/AffineStructures.h"
+#include "mlir/Dialect/Linalg/Transforms/ValueBoundsOpInterface.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+
+using namespace mlir;
+using namespace mlir::linalg;
+using presburger::IntegerPolyhedron;
+
+namespace mlir {
+namespace tensor {
+namespace {
+
+struct CastOpInterface
+    : public ValueBoundsOpInterface::ExternalModel<CastOpInterface, CastOp> {
+  void populateBoundsForShapedValueDim(Operation *op, Value value, int64_t dim,
+                                       ValueBoundsConstraintSet &cstr) const {
+    auto castOp = cast<CastOp>(op);
+    assert(value == castOp.getResult() && "invalid value");
+
+    if (castOp.getResult().getType().isa<RankedTensorType>() &&
+        castOp.getSource().getType().isa<RankedTensorType>()) {
+      cstr.addBound(IntegerPolyhedron::BoundType::EQ, value, dim,
+                    cstr.getExpr(castOp.getSource(), dim));
+    }
+  }
+};
+
+struct DimOpInterface
+    : public ValueBoundsOpInterface::ExternalModel<DimOpInterface, DimOp> {
+  void populateBoundsForIndexValue(Operation *op, Value value,
+                                   ValueBoundsConstraintSet &cstr) const {
+    auto dimOp = cast<DimOp>(op);
+    assert(value == dimOp.getResult() && "invalid value");
+
+    auto constIndex = dimOp.getConstantIndex();
+    if (!constIndex.has_value())
+      return;
+    cstr.addBound(IntegerPolyhedron::BoundType::EQ, value,
+                  cstr.getExpr(dimOp.getSource(), *constIndex));
+  }
+};
+
+struct EmptyOpInterface
+    : public ValueBoundsOpInterface::ExternalModel<EmptyOpInterface, EmptyOp> {
+  void populateBoundsForShapedValueDim(Operation *op, Value value, int64_t dim,
+                                       ValueBoundsConstraintSet &cstr) const {
+    auto emptyOp = cast<EmptyOp>(op);
+    assert(value == emptyOp.getResult() && "invalid value");
+
+    cstr.addBound(IntegerPolyhedron::BoundType::EQ, value, dim,
+                  cstr.getExpr(emptyOp.getMixedSizes()[dim]));
+  }
+};
+
+struct ExtractSliceOpInterface
+    : public ValueBoundsOpInterface::ExternalModel<ExtractSliceOpInterface,
+                                                   ExtractSliceOp> {
+  void populateBoundsForShapedValueDim(Operation *op, Value value, int64_t dim,
+                                       ValueBoundsConstraintSet &cstr) const {
+    auto extractSliceOp = cast<ExtractSliceOp>(op);
+    assert(value == extractSliceOp.getResult() && "invalid value");
+
+    llvm::SmallBitVector dropped = extractSliceOp.getDroppedDims();
+    int64_t ctr = -1;
+    for (int64_t i = 0, e = extractSliceOp.getMixedSizes().size(); i < e; ++i) {
+      // Skip over rank-reduced dimensions.
+      if (!dropped.test(i))
+        ++ctr;
+      if (ctr == dim) {
+        cstr.addBound(IntegerPolyhedron::BoundType::EQ, value, dim,
+                      extractSliceOp.getMixedSizes()[i]);
+        return;
+      }
+    }
+    llvm_unreachable("could not find non-rank-reduced dim");
+  }
+};
+
+struct PadOpInterface
+    : public ValueBoundsOpInterface::ExternalModel<PadOpInterface, PadOp> {
+  void populateBoundsForShapedValueDim(Operation *op, Value value, int64_t dim,
+                                       ValueBoundsConstraintSet &cstr) const {
+    auto padOp = cast<PadOp>(op);
+    assert(value == padOp.getResult() && "invalid value");
+
+    AffineExpr expr = cstr.getExpr(padOp.getSource(), dim) +
+                      cstr.getExpr(padOp.getMixedLowPad()[dim]) +
+                      cstr.getExpr(padOp.getMixedHighPad()[dim]);
+    cstr.addBound(IntegerPolyhedron::BoundType::EQ, value, dim, expr);
+  }
+};
+
+struct RankOpInterface
+    : public ValueBoundsOpInterface::ExternalModel<RankOpInterface, RankOp> {
+  void populateBoundsForIndexValue(Operation *op, Value value,
+                                   ValueBoundsConstraintSet &cstr) const {
+    auto rankOp = cast<RankOp>(op);
+    assert(value == rankOp.getResult() && "invalid value");
+
+    auto tensorType = rankOp.getTensor().getType().dyn_cast<RankedTensorType>();
+    if (!tensorType)
+      return;
+    cstr.addBound(IntegerPolyhedron::BoundType::EQ, value,
+                  cstr.getExpr(tensorType.getRank()));
+  }
+};
+
+} // namespace
+} // namespace tensor
+} // namespace mlir
+
+void mlir::linalg::registerValueBoundsOpInterfaceExternalModels(
+    DialectRegistry &registry) {
+  registry.addExtension(+[](MLIRContext *ctx, tensor::TensorDialect *dialect) {
+    tensor::CastOp::attachInterface<tensor::CastOpInterface>(*ctx);
+    tensor::DimOp::attachInterface<tensor::DimOpInterface>(*ctx);
+    tensor::EmptyOp::attachInterface<tensor::EmptyOpInterface>(*ctx);
+    tensor::ExtractSliceOp::attachInterface<tensor::ExtractSliceOpInterface>(
+        *ctx);
+    tensor::InsertOp::attachInterface<
+        DstValueBoundsOpInterfaceExternalModel<tensor::InsertOp>>(*ctx);
+    tensor::InsertSliceOp::attachInterface<
+        DstValueBoundsOpInterfaceExternalModel<tensor::InsertSliceOp>>(*ctx);
+    tensor::PadOp::attachInterface<tensor::PadOpInterface>(*ctx);
+    tensor::RankOp::attachInterface<tensor::RankOpInterface>(*ctx);
+  });
+}

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

+// RUN: mlir-opt %s -test-linalg-transform-patterns=test-refiy-shape-dims \
+// RUN:     -verify-diagnostics -split-input-file -allow-unregistered-dialect | FileCheck %s
+
+func.func @unknown_op() -> index {
+  %0 = "test.foo"() : () -> (tensor<?x?xf32>)
+  // expected-error @below{{could not reify bound}}
+  %1 = "test.reify_bound"(%0) {dim = 0} : (tensor<?x?xf32>) -> (index)
+  return %1 : index
+}
+
+// -----
+
+// CHECK-LABEL: func @cast(
+//       CHECK:   %[[c10:.*]] = arith.constant 10 : index
+//       CHECK:   return %[[c10]]
+func.func @cast(%t: tensor<10xf32>) -> index {
+  %0 = tensor.cast %t : tensor<10xf32> to tensor<?xf32>
+  %1 = "test.reify_bound"(%0) {dim = 0} : (tensor<?xf32>) -> (index)
+  return %1 : index
+}
+
+// -----
+
+func.func @cast_unranked(%t: tensor<*xf32>) -> index {
+  %0 = tensor.cast %t : tensor<*xf32> to tensor<?xf32>
+  // expected-error @below{{could not reify bound}}
+  %1 = "test.reify_bound"(%0) {dim = 0} : (tensor<?xf32>) -> (index)
+  return %1 : index
+}
+
+// -----
+
+// CHECK-LABEL: func @dim(
+//  CHECK-SAME:     %[[t:.*]]: tensor<?xf32>
+//       CHECK:   %[[dim:.*]] = tensor.dim %[[t]]
+//       CHECK:   %[[dim:.*]] = tensor.dim %[[t]]
+//       CHECK:   return %[[dim]]
+func.func @dim(%t: tensor<?xf32>) -> index {
+  %c0 = arith.constant 0 : index
+  %0 = tensor.dim %t, %c0 : tensor<?xf32>
+  %1 = "test.reify_bound"(%0) : (index) -> (index)
+  return %1 : index
+}
+
+// -----
+
+// CHECK-LABEL: func @empty(
+//  CHECK-SAME:     %[[sz:.*]]: index
+//       CHECK:   %[[c6:.*]] = arith.constant 6 : index
+//       CHECK:   return %[[c6]], %[[sz]]
+func.func @empty(%sz: index) -> (index, index) {
+  %0 = tensor.empty(%sz) : tensor<6x?xf32>
+  %1 = "test.reify_bound"(%0) {dim = 0} : (tensor<6x?xf32>) -> (index)
+  %2 = "test.reify_bound"(%0) {dim = 1} : (tensor<6x?xf32>) -> (index)
+  return %1, %2 : index, index
+}
+
+// -----
+
+// CHECK-LABEL: func @extract_slice_dynamic(
+//  CHECK-SAME:     %[[t:.*]]: tensor<?xf32>, %[[sz:.*]]: index
+//       CHECK:   return %[[sz]]
+func.func @extract_slice_dynamic(%t: tensor<?xf32>, %sz: index) -> index {
+  %0 = tensor.extract_slice %t[2][%sz][1] : tensor<?xf32> to tensor<?xf32>
+  %1 = "test.reify_bound"(%0) {dim = 0} : (tensor<?xf32>) -> (index)
+  return %1 : index
+}
+
+// -----
+
+// CHECK-LABEL: func @extract_slice_static(
+//  CHECK-SAME:     %[[t:.*]]: tensor<?xf32>
+//       CHECK:   %[[c5:.*]] = arith.constant 5 : index
+//       CHECK:   return %[[c5]]
+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
+}
+
+// -----
+
+// 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
+}
+
+// -----
+
+// CHECK-LABEL: func @insert(
+//  CHECK-SAME:     %[[t:.*]]: tensor<?xf32>
+//       CHECK:   %[[c0:.*]] = arith.constant 0 : index
+//       CHECK:   %[[dim:.*]] = tensor.dim %[[t]], %[[c0]]
+//       CHECK:   return %[[dim]]
+func.func @insert(%t: tensor<?xf32>, %f: f32, %pos: index) -> index {
+  %0 = tensor.insert %f into %t[%pos] : tensor<?xf32>
+  %1 = "test.reify_bound"(%0) {dim = 0} : (tensor<?xf32>) -> (index)
+  return %1 : index
+}
+
+// -----
+
+// CHECK: #[[$map:.*]] = affine_map<()[s0] -> (s0 + 12)>
+// CHECK: #[[$map1:.*]] = affine_map<()[s0, s1] -> (s0 + s1 * 2)>
+// CHECK-LABEL: func @pad(
+//  CHECK-SAME:     %[[t:.*]]: tensor<?x7xf32>, %[[a:.*]]: index, %[[b:.*]]: index
+//       CHECK:   %[[bound1:.*]] = affine.apply #[[$map]]()[%[[b]]]
+//       CHECK:   %[[c0:.*]] = arith.constant 0 : index
+//       CHECK:   %[[dim0:.*]] = tensor.dim %[[t]], %[[c0]]
+//       CHECK:   %[[bound0:.*]] = affine.apply #[[$map1]]()[%[[dim0]], %[[a]]]
+//       CHECK:   return %[[bound0]], %[[bound1]]
+func.func @pad(%t: tensor<?x7xf32>, %a: index, %b: index) -> (index, index) {
+  %pad = arith.constant 0.0 : f32
+  %0 = tensor.pad %t low[%a, 5] high[%a, %b] {
+    ^bb0(%arg1: index, %arg2: index):
+      tensor.yield %pad : f32
+    } : tensor<?x7xf32> to tensor<?x?xf32>
+  %1 = "test.reify_bound"(%0) {dim = 0} : (tensor<?x?xf32>) -> (index)
+  %2 = "test.reify_bound"(%0) {dim = 1} : (tensor<?x?xf32>) -> (index)
+  return %1, %2 : index, index
+}
+
+// -----
+
+// CHECK-LABEL: func @rank(
+//  CHECK-SAME:     %[[t:.*]]: tensor<5xf32>
+//       CHECK:   %[[c1:.*]] = arith.constant 1 : index
+//       CHECK:   return %[[c1]]
+func.func @rank(%t: tensor<5xf32>) -> index {
+  %0 = tensor.rank %t : tensor<5xf32>
+  %1 = "test.reify_bound"(%0) : (index) -> (index)
+  return %1 : index
+}

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

+// RUN: mlir-opt %s -test-linalg-transform-patterns="test-refiy-shape-dims reify-to-func-args" \
+// RUN:     -verify-diagnostics -split-input-file -allow-unregistered-dialect | FileCheck %s
+
+// CHECK-LABEL: func @reify_through_chain(
+//  CHECK-SAME:     %[[sz0:.*]]: index, %[[sz2:.*]]: index
+//       CHECK:   %[[c10:.*]] = arith.constant 10 : index
+//       CHECK:   return %[[sz0]], %[[c10]], %[[sz2]]
+func.func @reify_through_chain(%sz0: index, %sz2: index) -> (index, index, index) {
+  %c2 = arith.constant 2 : index
+  %0 = tensor.empty(%sz0, %sz2) : tensor<?x10x?xf32>
+  %1 = tensor.cast %0 : tensor<?x10x?xf32> to tensor<?x?x?xf32>
+  %pos = arith.constant 0 : index
+  %f = arith.constant 0.0 : f32
+  %2 = tensor.insert %f into %1[%pos, %pos, %pos] : tensor<?x?x?xf32>
+  %3 = tensor.dim %2, %c2 : tensor<?x?x?xf32>
+
+  %4 = "test.reify_bound"(%2) {dim = 0} : (tensor<?x?x?xf32>) -> (index)
+  %5 = "test.reify_bound"(%2) {dim = 1} : (tensor<?x?x?xf32>) -> (index)
+  %6 = "test.reify_bound"(%3) : (index) -> (index)
+
+  return %4, %5, %6 : index, index, index
+}

mlir/test/lib/Dialect/Linalg/TestLinalgTransforms.cpp

 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/Passes.h"
 #include "mlir/Dialect/Linalg/Transforms/Hoisting.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
+#include "mlir/Dialect/Linalg/Transforms/ValueBoundsOpInterface.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
 #include "llvm/ADT/SmallVector.h"
 
 using namespace mlir;
struct TestLinalgTransforms
   Option<bool> testEraseUnusedOperandsAndResults{
       *this, "test-erase-unused-operands-and-results",
       llvm::cl::desc("Test patterns to erase unused operands and results"),
       llvm::cl::init(false)};
   Option<bool> testEraseUnnecessaryInputs{
       *this, "test-erase-unnecessary-inputs",
       llvm::cl::desc("Test patterns to erase unnecessary inputs"),
       llvm::cl::init(false)};
+  Option<bool> testReifyBounds{*this, "test-refiy-shape-dims",
+                               llvm::cl::desc("Test value bounds reification"),
+                               llvm::cl::init(false)};
+  Option<bool> reifyToFuncArgs{
+      *this, "reify-to-func-args",
+      llvm::cl::desc("Reify in terms of function args"), llvm::cl::init(false)};
 };
 } // namespace
 
 static void applyPatterns(func::FuncOp funcOp) {
   MLIRContext *ctx = funcOp.getContext();
   RewritePatternSet patterns(ctx);
 
   //===--------------------------------------------------------------------===//

 }
 
 static void applyEraseUnnecessaryInputs(func::FuncOp funcOp) {
   RewritePatternSet patterns(funcOp.getContext());
   populateEraseUnnecessaryInputsPatterns(patterns);
   (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
 }
 
+static LogicalResult reifyValueBounds(func::FuncOp funcOp,
+                                      bool reifyToFuncArgs) {
+  IRRewriter rewriter(funcOp.getContext());
+  WalkResult result = funcOp.walk([&](Operation *op) {
+    if (op->getName().getStringRef() == "test.reify_bound") {
+      if (op->getNumOperands() != 1 || op->getNumResults() != 1 ||
+          !op->getResultTypes()[0].isIndex()) {
+        op->emitOpError("invalid op");
+        return WalkResult::skip();
+      }
+      Value value = op->getOperand(0);
+      if (value.getType().isa<IndexType>() !=
+          !op->hasAttrOfType<IntegerAttr>("dim")) {
+        // Op should have "dim" attribute if and only if the operand is an
+        // index-typed value.
+        op->emitOpError("invalid op");
+        return WalkResult::skip();
+      }
+      int64_t dim = value.getType().isIndex()
+                        ? ValueBoundsConstraintSet::kIndexValue
+                        : op->getAttrOfType<IntegerAttr>("dim").getInt();
+
+      rewriter.setInsertionPointAfterValue(value);
+      FailureOr<OpFoldResult> reified;
+      if (!reifyToFuncArgs) {
+        reified = ValueBoundsConstraintSet::reifyBound(
+            rewriter, op->getLoc(),
+            presburger::IntegerPolyhedron::BoundType::EQ, value, dim);
+      } else {
+        auto stopCondition = [](Value v) {
+          auto bbArg = v.dyn_cast<BlockArgument>();
+          if (!bbArg)
+            return false;
+          return isa<FunctionOpInterface>(
+              bbArg.getParentBlock()->getParentOp());
+        };
+        reified = ValueBoundsConstraintSet::reifyBound(
+            rewriter, op->getLoc(),
+            presburger::IntegerPolyhedron::BoundType::EQ, value, dim,
+            stopCondition);
+      }
+      if (failed(reified)) {
+        op->emitOpError("could not reify bound");
+        return WalkResult::interrupt();
+      }
+      if (auto val = reified->dyn_cast<Value>()) {
+        rewriter.replaceOp(op, val);
+        return WalkResult::skip();
+      }
+      Value constOp = rewriter.create<arith::ConstantIndexOp>(
+          op->getLoc(), reified->get<Attribute>().cast<IntegerAttr>().getInt());
+      rewriter.replaceOp(op, constOp);
+      return WalkResult::skip();
+    }
+    return WalkResult::advance();
+  });
+  return failure(result.wasInterrupted());
+}
+
 /// Apply transformations specified as patterns.
 void TestLinalgTransforms::runOnOperation() {
   if (testPatterns)
     return applyPatterns(getOperation());
   if (testVectorTransferForwardingPatterns)
     return applyVectorTransferForwardingPatterns(getOperation());
   if (testGenericToVectorPattern)
     return applyLinalgToVectorPatterns(getOperation());
void TestLinalgTransforms::runOnOperation() {
   if (testBubbleUpExtractSliceOpPattern)
     return applyBubbleUpExtractSliceOpPattern(getOperation());
   if (testSwapExtractSliceWithFill)
     return applySwapExtractSliceWithFillPattern(getOperation());
   if (testEraseUnusedOperandsAndResults)
     return applyEraseUnusedOperandsAndResultsPatterns(getOperation());
   if (testEraseUnnecessaryInputs)
     return applyEraseUnnecessaryInputs(getOperation());
+  if (testReifyBounds)
+    if (failed(reifyValueBounds(getOperation(), reifyToFuncArgs)))
+      return signalPassFailure();
 }
 
 namespace mlir {
 namespace test {
 void registerTestLinalgTransforms() {
   PassRegistration<TestLinalgTransforms>();
 }
 } // namespace test
 } // namespace mlir

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

deps = [
         ":Support",
         ":TensorDialect",
         ":TensorTilingInterfaceImpl",
         ":TensorTransforms",
         ":TensorUtils",
         ":TilingInterface",
         ":TransformUtils",
         ":Transforms",
+        ":ValueBoundsOpInterfaceIncGen",
         ":VectorDialect",
         ":VectorToSCF",
         ":VectorTransforms",
         ":VectorUtils",
         ":X86VectorTransforms",
         "//llvm:Support",
     ],
 )
 
+td_library(
+    name = "ValueBoundsOpInterfaceTdFiles",
+    srcs = [
+        "include/mlir/Dialect/Linalg/Transforms/ValueBoundsOpInterface.td",
+    ],
+    includes = ["include"],
+    deps = [
+        ":OpBaseTdFiles",
+    ],
+)
+
+gentbl_cc_library(
+    name = "ValueBoundsOpInterfaceIncGen",
+    strip_include_prefix = "include",
+    tbl_outs = [
+        (
+            ["-gen-op-interface-decls"],
+            "include/mlir/Dialect/Linalg/Transforms/ValueBoundsOpInterface.h.inc",
+        ),
+        (
+            ["-gen-op-interface-defs"],
+            "include/mlir/Dialect/Linalg/Transforms/ValueBoundsOpInterface.cpp.inc",
+        ),
+    ],
+    tblgen = ":mlir-tblgen",
+    td_file = "include/mlir/Dialect/Linalg/Transforms/ValueBoundsOpInterface.td",
+    deps = [
+        ":ValueBoundsOpInterfaceTdFiles",
+    ],
+)
+
 cc_library(
     name = "TilingInterface",
     srcs = ["lib/Interfaces/TilingInterface.cpp"],
     hdrs = ["include/mlir/Interfaces/TilingInterface.h"],
     includes = ["include"],
     deps = [
         ":DialectUtils",
         ":IR",