[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 ↗

(On Diff #503721)

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 ↗

(On Diff #503721)

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

27 ↗

(On Diff #503721)

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 ↗

(On Diff #503721)

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 ↗

(On Diff #503721)

(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 ↗

(On Diff #503721)

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 ↗

(On Diff #503721)

typo

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

228 ↗

(On Diff #503721)

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 ↗

(On Diff #503721)

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 ↗

(On Diff #503721)

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 ↗

(On Diff #503721)

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 ↗

(On Diff #503721)

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 ↗

(On Diff #503721)

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

27 ↗

(On Diff #503721)

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 ↗

(On Diff #503721)

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

14 ↗

(On Diff #503721)

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 ↗

(On Diff #503721)

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 ↗

(On Diff #503721)

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 ↗

(On Diff #503721)

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 ↗

(On Diff #503721)

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 ↗

(On Diff #503721)

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 ↗

(On Diff #503721)

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 ↗

(On Diff #503721)

(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 ↗

(On Diff #503721)

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 ↗

(On Diff #503721)

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 ↗

(On Diff #503721)

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

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

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

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

mlir/include/mlir/Interfaces/ValueBoundsOpInterface.h

36

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

44

OOC are all BoundType's allowed?

61

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

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

91

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

101

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

147

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

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

23

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

53

New form is

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

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

57

(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

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

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

119

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

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

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

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

I added an assertion in ValueBoundsConstraintSet::computeBound.

59

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

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

44

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

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

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

91

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

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

What do you mean by "more declarative"?

119

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

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

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

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

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

50

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

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

30

Update description ?

[springerm]

Thanks for the thorough review!

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

52

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/Affine/Transforms/Transforms.h

 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_AFFINE_TRANSFORMS_TRANSFORMS_H
 #define MLIR_DIALECT_AFFINE_TRANSFORMS_TRANSFORMS_H
 
 #include "mlir/Support/LogicalResult.h"
 
 namespace mlir {
+class AffineApplyOp;
+class Location;
+class OpBuilder;
+class OpFoldResult;
 class RewritePatternSet;
 class RewriterBase;
-class AffineApplyOp;
+class Value;
+
+namespace presburger {
+enum class BoundType;
+} // namespace presburger
 
 /// Populate patterns that expand affine index operations into more fundamental
 /// operations (not necessarily restricted to Affine dialect).
 void populateAffineExpandIndexOpsPatterns(RewritePatternSet &patterns);
 
 /// Helper function to rewrite `op`'s affine map and reorder its operands such
 /// that they are in increasing order of hoistability (i.e. the least hoistable)
 /// operands come first in the operand list.
 void reorderOperandsByHoistability(RewriterBase &rewriter, AffineApplyOp op);
 
 /// Split an "affine.apply" operation into smaller ops.
 /// This reassociates a large AffineApplyOp into an ordered list of smaller
 /// AffineApplyOps. This can be used right before lowering affine ops to arith
 /// to exhibit more opportunities for CSE and LICM.
 /// Return the sink AffineApplyOp on success or failure if `op` does not
 /// decompose into smaller AffineApplyOps.
 /// Note that this can be undone by canonicalization which tries to
 /// maximally compose chains of AffineApplyOps.
 FailureOr<AffineApplyOp> decompose(RewriterBase &rewriter, AffineApplyOp op);
 
+/// Reify a bound for the given index-typed value or shape dimension size in
+/// terms of the owning op's operands. `dim` must be `nullopt` if and only if
+/// `value` is index-typed.
+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.
+FailureOr<OpFoldResult>
+reifyValueBound(OpBuilder &b, Location loc, presburger::BoundType type,
+                Value value, std::optional<int64_t> dim,
+                function_ref<bool(Value)> stopCondition);
+
 } // namespace mlir
 
 #endif // MLIR_DIALECT_AFFINE_TRANSFORMS_TRANSFORMS_H

mlir/include/mlir/Dialect/Tensor/IR/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_TENSOR_IR_VALUEBOUNDSOPINTERFACEIMPL_H
+#define MLIR_DIALECT_TENSOR_IR_VALUEBOUNDSOPINTERFACEIMPL_H
+
+namespace mlir {
+class DialectRegistry;
+
+namespace tensor {
+void registerValueBoundsOpInterfaceExternalModels(DialectRegistry &registry);
+} // namespace tensor
+} // namespace mlir
+
+#endif // MLIR_DIALECT_TENSOR_IR_VALUEBOUNDSOPINTERFACEIMPL_H

mlir/include/mlir/InitAllDialects.h

 #include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
 #include "mlir/Dialect/Shape/IR/Shape.h"
 #include "mlir/Dialect/Shape/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
 #include "mlir/Dialect/SparseTensor/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Tensor/IR/TensorInferTypeOpInterfaceImpl.h"
 #include "mlir/Dialect/Tensor/IR/TensorTilingInterfaceImpl.h"
+#include "mlir/Dialect/Tensor/IR/ValueBoundsOpInterfaceImpl.h"
 #include "mlir/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/Tosa/IR/TosaOps.h"
 #include "mlir/Dialect/Transform/IR/TransformDialect.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/Dialect/Vector/TransformOps/VectorTransformOps.h"
 #include "mlir/Dialect/Vector/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/X86Vector/X86VectorDialect.h"
 #include "mlir/IR/Dialect.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);
   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);
+  tensor::registerValueBoundsOpInterfaceExternalModels(registry);
   vector::registerBufferizableOpInterfaceExternalModels(registry);
 }
 
 /// Append all the MLIR dialects to the registry contained in the given context.
 inline void registerAllDialects(MLIRContext &context) {
   DialectRegistry registry;
   registerAllDialects(registry);
   context.appendDialectRegistry(registry);
 }
 
 } // namespace mlir
 
 #endif // MLIR_INITALLDIALECTS_H_

mlir/include/mlir/Interfaces/CMakeLists.txt

 add_mlir_interface(CallInterfaces)
 add_mlir_interface(CastInterfaces)
 add_mlir_interface(ControlFlowInterfaces)
 add_mlir_interface(CopyOpInterface)
 add_mlir_interface(DerivedAttributeOpInterface)
 add_mlir_interface(DestinationStyleOpInterface)
 add_mlir_interface(InferIntRangeInterface)
 add_mlir_interface(InferTypeOpInterface)
 add_mlir_interface(LoopLikeInterface)
 add_mlir_interface(ParallelCombiningOpInterface)
 add_mlir_interface(RuntimeVerifiableOpInterface)
 add_mlir_interface(ShapedOpInterfaces)
 add_mlir_interface(SideEffectInterfaces)
 add_mlir_interface(TilingInterface)
+add_mlir_interface(ValueBoundsOpInterface)
 add_mlir_interface(VectorInterfaces)
 add_mlir_interface(ViewLikeInterface)
 
 set(LLVM_TARGET_DEFINITIONS DataLayoutInterfaces.td)
 mlir_tablegen(DataLayoutAttrInterface.h.inc -gen-attr-interface-decls)
 mlir_tablegen(DataLayoutAttrInterface.cpp.inc -gen-attr-interface-defs)
 mlir_tablegen(DataLayoutOpInterface.h.inc -gen-op-interface-decls)
 mlir_tablegen(DataLayoutOpInterface.cpp.inc -gen-op-interface-defs)

mlir/include/mlir/Interfaces/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_INTERFACES_VALUEBOUNDSOPINTERFACE_H_
+#define MLIR_INTERFACES_VALUEBOUNDSOPINTERFACE_H_
+
+#include "mlir/Analysis/FlatLinearValueConstraints.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 {
+
+using ValueDimList = SmallVector<std::pair<Value, std::optional<int64_t>>>;
+
+/// A helper class to be used with `ValueBoundsOpInterface`. This class stores a
+/// constraint system and mapping of constrained variables to index-typed
+/// values or dimension sizes of shaped values.
+///
+/// Interface implementations of `ValueBoundsOpInterface` use `addBounds` to
+/// insert constraints about their results and/or region block arguments into
+/// the constraint set in the form of an AffineExpr. When a bound should be
+/// 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:
+  /// 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)>;
+
+  /// 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.
+  static LogicalResult computeBound(AffineMap &resultMap,
+                                    ValueDimList &mapOperands,
+                                    presburger::BoundType type, Value value,
+                                    std::optional<int64_t> dim,
+                                    StopConditionFn stopCondition);
+
+  /// Bound the given index-typed value by the given expression.
+  void addBound(presburger::BoundType type, Value value, AffineExpr expr);
+
+  /// Bound the given shaped value dimension by the given expression.
+  void addBound(presburger::BoundType type, Value value, int64_t dim,
+                AffineExpr expr);
+
+  /// Bound the given index-typed value by a constant or SSA value.
+  void addBound(presburger::BoundType type, Value value, OpFoldResult ofr);
+
+  /// Bound the given shaped value dimension by a constant or SSA value.
+  void addBound(presburger::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.
+  ///
+  /// `dim` must be `nullopt` if and only if the given value is of index type.
+  AffineExpr getExpr(Value value, std::optional<int64_t> dim = std::nullopt);
+
+  /// 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 constant);
+
+protected:
+  /// Dimension identifier to indicate a value is index-typed.
+  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(StopConditionFn 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;
+
+  /// 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);
+
+  /// Project out all columns for which the condition holds.
+  void projectOut(function_ref<bool(ValueDim)> condition);
+
+  /// Mapping of columns to values/shape dimensions.
+  SmallVector<ValueDim> positionToValueDim;
+  /// Reverse mapping of values/shape dimensions to columns.
+  DenseMap<ValueDim, int64_t> valueDimToPosition;
+
+  /// Worklist of values/shape dimensions that have not been processed yet.
+  SetVector<int64_t> worklist;
+
+  /// Constraint system of equalities and inequalities.
+  FlatLinearConstraints cstr;
+
+  /// Builder for constructing affine expressions.
+  Builder builder;
+};
+
+} // namespace mlir
+
+#include "mlir/Interfaces/ValueBoundsOpInterface.h.inc"
+
+namespace mlir {
+
+/// 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::BoundType::EQ, value, dim,
+                  cstr.getExpr(tiedOperand, dim));
+  }
+};
+
+} // namespace mlir
+
+#endif // MLIR_INTERFACES_VALUEBOUNDSOPINTERFACE_H_

mlir/include/mlir/Interfaces/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";
+  let methods = [
+    InterfaceMethod<
+      /*desc=*/[{
+        Populate the constraint set with bounds for the given index-typed
+        value.
+
+        Note: If `value` is a block argument, it must belong to an entry block
+        of a region. Unstructured control flow graphs are not supported at the
+        moment.
+      }],
+      /*retType=*/"void",
+      /*methodName=*/"populateBoundsForIndexValue",
+      /*args=*/(ins "::mlir::Value":$value,
+                    "::mlir::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.
+
+        Note: If `value` is a block argument, it must belong to an entry block
+        of a region. Unstructured control flow graphs are not supported at the
+        moment.
+      }],
+      /*retType=*/"void",
+      /*methodName=*/"populateBoundsForShapedValueDim",
+      /*args=*/(ins "::mlir::Value":$value,
+                    "int64_t":$dim,
+                    "::mlir::ValueBoundsConstraintSet &":$cstr),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        llvm_unreachable("not implemented");
+      }]
+    >,
+  ];
+}
+
+#endif  // VALUEBOUNDSOPINTERFACE

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

 add_mlir_dialect_library(MLIRAffineTransforms
   AffineDataCopyGeneration.cpp
   AffineExpandIndexOps.cpp
   AffineLoopInvariantCodeMotion.cpp
   AffineLoopNormalize.cpp
   AffineParallelize.cpp
   AffineScalarReplacement.cpp
   DecomposeAffineOps.cpp
   LoopCoalescing.cpp
   LoopFusion.cpp
   LoopTiling.cpp
   LoopUnroll.cpp
   LoopUnrollAndJam.cpp
   PipelineDataTransfer.cpp
+  ReifyValueBounds.cpp
   SuperVectorize.cpp
   SimplifyAffineStructures.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Affine
 
   DEPENDS
   MLIRAffineOpsIncGen
   MLIRAffinePassIncGen
   MLIRLoopLikeInterfaceIncGen
 
   LINK_LIBS PUBLIC
   MLIRAffineDialect
   MLIRAffineAnalysis
   MLIRAffineUtils
   MLIRArithDialect
   MLIRIR
   MLIRMemRefDialect
   MLIRPass
   MLIRSCFUtils
   MLIRSideEffectInterfaces
+  MLIRTensorDialect
   MLIRTransformUtils
+  MLIRValueBoundsOpInterface
   MLIRVectorDialect
   MLIRVectorUtils
   MLIRVectorToLLVM
   )
 

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

+//===- ReifyValueBounds.cpp --- Reify value bounds with affine ops ------*-===//
+//
+// 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/Transforms/Transforms.h"
+
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#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) {
+    // 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) {
+  // Compute bound.
+  AffineMap boundMap;
+  ValueDimList mapOperands;
+  if (failed(ValueBoundsConstraintSet::computeBound(boundMap, mapOperands, type,
+                                                    value, dim, stopCondition)))
+    return failure();
+
+  // Materialize tensor.dim/memref.dim ops.
+  SmallVector<Value> operands;
+  for (auto valueDim : mapOperands) {
+    Value value = valueDim.first;
+    std::optional<int64_t> dim = valueDim.second;
+
+    if (!dim.has_value()) {
+      // This is an index-typed value.
+      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");
+    }
+  }
+
+  // Simplify and return bound.
+  mlir::canonicalizeMapAndOperands(&boundMap, &operands);
+  // Check for special cases where no affine.apply op is needed.
+  if (boundMap.isSingleConstant()) {
+    // Bound is a constant: return an IntegerAttr.
+    return static_cast<OpFoldResult>(
+        b.getIndexAttr(boundMap.getSingleConstantResult()));
+  }
+  // No affine.apply op is needed if the bound is a single SSA value.
+  if (auto expr = boundMap.getResult(0).dyn_cast<AffineDimExpr>())
+    return static_cast<OpFoldResult>(operands[expr.getPosition()]);
+  if (auto expr = boundMap.getResult(0).dyn_cast<AffineSymbolExpr>())
+    return static_cast<OpFoldResult>(
+        operands[expr.getPosition() + boundMap.getNumDims()]);
+  // General case: build affine.apply op.
+  return static_cast<OpFoldResult>(
+      b.create<AffineApplyOp>(loc, boundMap, operands).getResult());
+}

mlir/lib/Dialect/Tensor/IR/CMakeLists.txt

 set(LLVM_OPTIONAL_SOURCES
   TensorDialect.cpp
   TensorInferTypeOpInterfaceImpl.cpp
   TensorOps.cpp
   TensorTilingInterfaceImpl.cpp
+  ValueBoundsOpInterfaceImpl.cpp
 )
 
 add_mlir_dialect_library(MLIRTensorDialect
   TensorDialect.cpp
   TensorOps.cpp
+  ValueBoundsOpInterfaceImpl.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${PROJECT_SOURCE_DIR}/include/mlir/Dialect/Tensor
 
   DEPENDS
   MLIRTensorOpsIncGen
 
   LINK_COMPONENTS
   Core
 
   LINK_LIBS PUBLIC
   MLIRAffineDialect
   MLIRArithDialect
   MLIRArithUtils
   MLIRCastInterfaces
   MLIRComplexDialect
   MLIRDestinationStyleOpInterface
   MLIRDialectUtils
   MLIRIR
   MLIRInferTypeOpInterface
   MLIRParallelCombiningOpInterface
   MLIRShapedOpInterfaces
   MLIRSideEffectInterfaces
   MLIRSupport
+  MLIRValueBoundsOpInterface
   MLIRViewLikeInterface
   )
 
 add_mlir_dialect_library(MLIRTensorInferTypeOpInterfaceImpl
   TensorInferTypeOpInterfaceImpl.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Tensor

mlir/lib/Dialect/Tensor/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/Tensor/IR/ValueBoundsOpInterfaceImpl.h"
+
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Interfaces/ValueBoundsOpInterface.h"
+
+using namespace mlir;
+using presburger::BoundType;
+
+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(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(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(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(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(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(BoundType::EQ, value, cstr.getExpr(tensorType.getRank()));
+  }
+};
+
+} // namespace
+} // namespace tensor
+} // namespace mlir
+
+void mlir::tensor::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/lib/Interfaces/CMakeLists.txt

 set(LLVM_OPTIONAL_SOURCES
   CallInterfaces.cpp
   CastInterfaces.cpp
   ControlFlowInterfaces.cpp
   CopyOpInterface.cpp
   DataLayoutInterfaces.cpp
   DerivedAttributeOpInterface.cpp
   DestinationStyleOpInterface.cpp
   InferIntRangeInterface.cpp
   InferTypeOpInterface.cpp
   LoopLikeInterface.cpp
   ParallelCombiningOpInterface.cpp
   RuntimeVerifiableOpInterface.cpp
   ShapedOpInterfaces.cpp
   SideEffectInterfaces.cpp
   TilingInterface.cpp
+  ValueBoundsOpInterface.cpp
   VectorInterfaces.cpp
   ViewLikeInterface.cpp
   )
 
 function(add_mlir_interface_library name)
   add_mlir_library(MLIR${name}
     ${name}.cpp
 

 add_mlir_interface_library(InferIntRangeInterface)
 add_mlir_interface_library(InferTypeOpInterface)
 add_mlir_interface_library(LoopLikeInterface)
 add_mlir_interface_library(ParallelCombiningOpInterface)
 add_mlir_interface_library(RuntimeVerifiableOpInterface)
 add_mlir_interface_library(ShapedOpInterfaces)
 add_mlir_interface_library(SideEffectInterfaces)
 add_mlir_interface_library(TilingInterface)
+add_mlir_interface_library(ValueBoundsOpInterface)
 add_mlir_interface_library(VectorInterfaces)
 add_mlir_interface_library(ViewLikeInterface)
 
 add_subdirectory(Utils)

mlir/lib/Interfaces/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/Interfaces/ValueBoundsOpInterface.h"
+
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Matchers.h"
+#include "llvm/ADT/APSInt.h"
+
+using namespace mlir;
+using presburger::BoundType;
+using presburger::VarKind;
+
+namespace mlir {
+#include "mlir/Interfaces/ValueBoundsOpInterface.cpp.inc"
+} // namespace mlir
+
+/// If ofr is a constant integer or an IntegerAttr, return the integer.
+static std::optional<int64_t> getConstantIntValue(OpFoldResult ofr) {
+  // Case 1: Check for Constant integer.
+  if (auto val = ofr.dyn_cast<Value>()) {
+    APSInt intVal;
+    if (matchPattern(val, m_ConstantInt(&intVal)))
+      return intVal.getSExtValue();
+    return std::nullopt;
+  }
+  // Case 2: Check for IntegerAttr.
+  Attribute attr = ofr.dyn_cast<Attribute>();
+  if (auto intAttr = attr.dyn_cast_or_null<IntegerAttr>())
+    return intAttr.getValue().getSExtValue();
+  return std::nullopt;
+}
+
+ValueBoundsConstraintSet::ValueBoundsConstraintSet(ValueDim valueDim)
+    : builder(valueDim.first.getContext()) {
+  insert(valueDim, /*isSymbol=*/false);
+}
+
+#ifndef NDEBUG
+static void assertValidValueDim(Value value, std::optional<int64_t> dim) {
+  if (value.getType().isIndex()) {
+    assert(!dim.has_value() && "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(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(BoundType type, Value value,
+                                        AffineExpr expr) {
+  assert(value.getType().isIndex() && "expected index type");
+  assert((value.isa<OpResult>() ||
+          value.cast<BlockArgument>().getOwner()->isEntryBlock()) &&
+         "unstructured control flow is not supported");
+
+  addBound(type, getPos(value), expr);
+}
+
+void ValueBoundsConstraintSet::addBound(BoundType type, Value value,
+                                        int64_t dim, AffineExpr expr) {
+#ifndef NDEBUG
+  assertValidValueDim(value, dim);
+  assert((value.isa<OpResult>() ||
+          value.cast<BlockArgument>().getOwner()->isEntryBlock()) &&
+         "unstructured control flow is not supported");
+#endif // NDEBUG
+
+  addBound(type, getPos(value, dim), expr);
+}
+
+void ValueBoundsConstraintSet::addBound(BoundType type, Value value,
+                                        OpFoldResult ofr) {
+  assert(value.getType().isIndex() && "expected index type");
+  addBound(type, getPos(value), getExpr(ofr));
+}
+
+void ValueBoundsConstraintSet::addBound(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,
+                                             std::optional<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.value_or(kIndexValue));
+  if (valueDimToPosition.find(valueDim) == valueDimToPosition.end())
+    (void)insert(valueDim);
+  int64_t pos = getPos(value, dim.value_or(kIndexValue));
+  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=*/std::nullopt);
+  auto constInt = getConstantIntValue(ofr);
+  assert(constInt.has_value() && "expected Integer constant");
+  return builder.getAffineConstantExpr(*constInt);
+}
+
+AffineExpr ValueBoundsConstraintSet::getExpr(int64_t constant) {
+  return builder.getAffineConstantExpr(constant);
+}
+
+int64_t ValueBoundsConstraintSet::insert(ValueDim valueDim, bool isSymbol) {
+  assert((valueDimToPosition.find(valueDim) == valueDimToPosition.end()) &&
+         "already mapped");
+  int64_t pos = isSymbol ? cstr.appendVar(VarKind::Symbol)
+                         : cstr.appendVar(VarKind::SetDim);
+  positionToValueDim.insert(positionToValueDim.begin() + pos, valueDim);
+  // Update reverse mapping.
+  for (int64_t i = pos; i < positionToValueDim.size(); ++i)
+    valueDimToPosition[positionToValueDim[i]] = i;
+
+  worklist.insert(pos);
+  return pos;
+}
+
+int64_t ValueBoundsConstraintSet::getPos(Value value, int64_t dim) const {
+#ifndef NDEBUG
+  assertValidValueDim(value, dim == kIndexValue
+                                 ? std::nullopt
+                                 : std::make_optional<int64_t>(dim));
+#endif // NDEBUG
+
+  auto it = valueDimToPosition.find(std::make_pair(value, dim));
+  assert(it != valueDimToPosition.end() && "expected mapped entry");
+  return it->second;
+}
+
+static Operation *getOwnerOfValue(Value value) {
+  if (auto bbArg = value.dyn_cast<BlockArgument>())
+    return bbArg.getOwner()->getParentOp();
+  return value.getDefiningOp();
+}
+
+void ValueBoundsConstraintSet::processWorklist(StopConditionFn 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(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) {
+  assert(pos >= 0 && pos < positionToValueDim.size() && "invalid position");
+  cstr.projectOut(pos);
+  bool erased = valueDimToPosition.erase(positionToValueDim[pos]);
+  (void)erased;
+  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;
+}
+
+void ValueBoundsConstraintSet::projectOut(
+    function_ref<bool(ValueDim)> condition) {
+  int64_t nextPos = 0;
+  while (nextPos < positionToValueDim.size()) {
+    if (condition(positionToValueDim[nextPos])) {
+      projectOut(nextPos);
+      // The column was projected out so another column is now at that position.
+      // Do not increase the counter.
+    } else {
+      ++nextPos;
+    }
+  }
+}
+
+LogicalResult ValueBoundsConstraintSet::computeBound(
+    AffineMap &resultMap, ValueDimList &mapOperands, presburger::BoundType type,
+    Value value, std::optional<int64_t> dim, StopConditionFn stopCondition) {
+  // Only EQ bounds are supported at the moment.
+  assert(type == BoundType::EQ && "unsupported bound type");
+
+  Builder b(value.getContext());
+  mapOperands.clear();
+
+  if (stopCondition(value)) {
+    // 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(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);
+  });
+
+  // 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);
+  // 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.
+  assert(cstr.cstr.getNumDimAndSymbolVars() == cstr.positionToValueDim.size() &&
+         "inconsistent mapping state");
+  SmallVector<AffineExpr> replacementDims, replacementSymbols;
+  int64_t numDims = 0, numSymbols = 0;
+  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;
+    bool isDim = i < cstr.cstr.getNumDimVars();
+    if (isDim) {
+      if (lb[0].isFunctionOfDim(i))
+        used = true;
+    } else {
+      if (lb[0].isFunctionOfSymbol(i - cstr.cstr.getNumDimVars()))
+        used = true;
+    }
+
+    if (!used) {
+      // Not used: Remove dim/symbol from the result.
+      if (isDim) {
+        replacementDims.push_back(b.getAffineConstantExpr(0));
+      } else {
+        replacementSymbols.push_back(b.getAffineConstantExpr(0));
+      }
+      continue;
+    }
+
+    if (isDim) {
+      replacementDims.push_back(b.getAffineDimExpr(numDims++));
+    } else {
+      replacementSymbols.push_back(b.getAffineSymbolExpr(numSymbols++));
+    }
+
+    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");
+      mapOperands.push_back(std::make_pair(value, std::nullopt));
+      continue;
+    }
+
+    assert(value.getType().cast<ShapedType>().isDynamicDim(dim) &&
+           "expected dynamic dim");
+    mapOperands.push_back(std::make_pair(value, dim));
+  }
+
+  resultMap = lb[0].replaceDimsAndSymbols(replacementDims, replacementSymbols,
+                                          numDims, numSymbols);
+  return success();
+}

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

+// RUN: mlir-opt %s -test-affine-reify-value-bounds="reify-to-func-args" \
+// RUN:     -verify-diagnostics -split-input-file | 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/Dialect/Tensor/value-bounds-op-interface-impl.mlir

+// RUN: mlir-opt %s -test-affine-reify-value-bounds -verify-diagnostics \
+// RUN:     -split-input-file | 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, s1] -> (s0 + s1 * 2)>
+// CHECK: #[[$map1:.*]] = affine_map<()[s0] -> (s0 + 12)>
+// CHECK-LABEL: func @pad(
+//  CHECK-SAME:     %[[t:.*]]: tensor<?x7xf32>, %[[a:.*]]: index, %[[b:.*]]: index
+//       CHECK:   %[[c0:.*]] = arith.constant 0 : index
+//       CHECK:   %[[dim0:.*]] = tensor.dim %[[t]], %[[c0]]
+//       CHECK:   %[[bound0:.*]] = affine.apply #[[$map]]()[%[[dim0]], %[[a]]]
+//       CHECK:   %[[bound1:.*]] = affine.apply #[[$map1]]()[%[[b]]]
+//       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/lib/Dialect/Affine/CMakeLists.txt

 # Exclude tests from libMLIR.so
 add_mlir_library(MLIRAffineTransformsTestPasses
   TestAffineDataCopy.cpp
   TestAffineLoopUnswitching.cpp
   TestAffineLoopParametricTiling.cpp
   TestDecomposeAffineOps.cpp
+  TestReifyValueBounds.cpp
   TestLoopFusion.cpp
   TestLoopMapping.cpp
   TestLoopPermutation.cpp
   TestVectorizationUtils.cpp
 
   EXCLUDE_FROM_LIBMLIR
 
   ADDITIONAL_HEADER_DIRS

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/IR/PatternMatch.h"
+#include "mlir/Interfaces/ValueBoundsOpInterface.h"
+#include "mlir/Pass/Pass.h"
+
+#define PASS_NAME "test-affine-reify-value-bounds"
+
+using namespace mlir;
+
+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 runOnOperation() override;
+
+private:
+  Option<bool> reifyToFuncArgs{
+      *this, "reify-to-func-args",
+      llvm::cl::desc("Reify in terms of function args"), llvm::cl::init(false)};
+};
+
+} // namespace
+
+/// Look for "test.reify_bound" ops in the input and replace their results with
+/// the reified values.
+static LogicalResult testReifyValueBounds(func::FuncOp funcOp,
+                                          bool reifyToFuncArgs) {
+  IRRewriter rewriter(funcOp.getContext());
+  WalkResult result = funcOp.walk([&](Operation *op) {
+    // Look for test.reify_bound ops.
+    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();
+      }
+
+      auto dim = value.getType().isIndex()
+                     ? std::nullopt
+                     : std::make_optional<int64_t>(
+                           op->getAttrOfType<IntegerAttr>("dim").getInt());
+
+      // Reify value bound.
+      rewriter.setInsertionPointAfter(op);
+      FailureOr<OpFoldResult> reified;
+      if (!reifyToFuncArgs) {
+        // Reify in terms of the op's operands.
+        reified = reifyValueBound(rewriter, op->getLoc(),
+                                  presburger::BoundType::EQ, value, dim);
+      } else {
+        // Reify in terms of function block arguments.
+        auto stopCondition = [](Value v) {
+          auto bbArg = v.dyn_cast<BlockArgument>();
+          if (!bbArg)
+            return false;
+          return isa<FunctionOpInterface>(
+              bbArg.getParentBlock()->getParentOp());
+        };
+        reified =
+            reifyValueBound(rewriter, op->getLoc(), presburger::BoundType::EQ,
+                            value, dim, stopCondition);
+      }
+      if (failed(reified)) {
+        op->emitOpError("could not reify bound");
+        return WalkResult::interrupt();
+      }
+
+      // Replace the op with the reified bound.
+      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());
+}
+
+void TestReifyValueBounds::runOnOperation() {
+  if (failed(testReifyValueBounds(getOperation(), reifyToFuncArgs)))
+    signalPassFailure();
+}
+
+namespace mlir {
+void registerTestAffineReifyValueBoundsPass() {
+  PassRegistration<TestReifyValueBounds>();
+}
+} // namespace mlir

mlir/tools/mlir-opt/mlir-opt.cpp

 void registerPrintSpirvAvailabilityPass();
 void registerLoopLikeInterfaceTestPasses();
 void registerShapeFunctionTestPasses();
 void registerSideEffectTestPasses();
 void registerSliceAnalysisTestPass();
 void registerSymbolTestPasses();
 void registerRegionTestPasses();
 void registerTestAffineDataCopyPass();
+void registerTestAffineReifyValueBoundsPass();
 void registerTestDecomposeAffineOpPass();
 void registerTestAffineLoopUnswitchingPass();
 void registerTestAllReduceLoweringPass();
 void registerTestFunc();
 void registerTestGpuMemoryPromotionPass();
 void registerTestLoopPermutationPass();
 void registerTestMatchers();
 void registerTestOperationEqualPass();
void registerTestPasses() {
   registerPrintSpirvAvailabilityPass();
   registerLoopLikeInterfaceTestPasses();
   registerShapeFunctionTestPasses();
   registerSideEffectTestPasses();
   registerSliceAnalysisTestPass();
   registerSymbolTestPasses();
   registerRegionTestPasses();
   registerTestAffineDataCopyPass();
+  registerTestAffineReifyValueBoundsPass();
   registerTestDecomposeAffineOpPass();
   registerTestAffineLoopUnswitchingPass();
   registerTestAllReduceLoweringPass();
   registerTestFunc();
   registerTestGpuMemoryPromotionPass();
   registerTestLoopPermutationPass();
   registerTestMatchers();
   registerTestOperationEqualPass();

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

deps = [
         ":ArithUtils",
         ":FuncDialect",
         ":IR",
         ":MemRefDialect",
         ":Pass",
         ":SCFDialect",
         ":SCFUtils",
         ":Support",
+        ":TensorDialect",
         ":Transforms",
+        ":ValueBoundsOpInterface",
         ":VectorDialect",
         ":VectorUtils",
         "//llvm:Support",
     ],
 )
 
 gentbl_cc_library(
     name = "ConversionPassIncGen",

 )
 
 cc_library(
     name = "TensorDialect",
     srcs = [
         "include/mlir/Transforms/InliningUtils.h",
         "lib/Dialect/Tensor/IR/TensorDialect.cpp",
         "lib/Dialect/Tensor/IR/TensorOps.cpp",
+        "lib/Dialect/Tensor/IR/ValueBoundsOpInterfaceImpl.cpp",
+    ],
+    hdrs = [
+        "include/mlir/Dialect/Tensor/IR/Tensor.h",
+        "include/mlir/Dialect/Tensor/IR/ValueBoundsOpInterfaceImpl.h",
     ],
-    hdrs = ["include/mlir/Dialect/Tensor/IR/Tensor.h"],
     includes = ["include"],
     deps = [
         ":AffineDialect",
         ":ArithDialect",
         ":ArithUtils",
         ":CastOpInterfaces",
         ":ComplexDialect",
         ":ControlFlowInterfaces",
         ":DestinationStyleOpInterface",
         ":DialectUtils",
         ":IR",
         ":InferTypeOpInterface",
         ":ParallelCombiningOpInterface",
         ":ShapedOpInterfaces",
         ":SideEffectInterfaces",
         ":Support",
         ":TensorOpsIncGen",
         ":TilingInterface",
+        ":ValueBoundsOpInterface",
         ":ViewLikeInterface",
         "//llvm:Support",
     ],
 )
 
 cc_library(
     name = "TensorInferTypeOpInterfaceImpl",
     srcs = ["lib/Dialect/Tensor/IR/TensorInferTypeOpInterfaceImpl.cpp"],
deps = [
         ":VectorToSCF",
         ":VectorTransforms",
         ":VectorUtils",
         ":X86VectorTransforms",
         "//llvm:Support",
     ],
 )
 
+td_library(
+    name = "ValueBoundsOpInterfaceTdFiles",
+    srcs = [
+        "include/mlir/Interfaces/ValueBoundsOpInterface.td",
+    ],
+    includes = ["include"],
+    deps = [
+        ":OpBaseTdFiles",
+    ],
+)
+
+gentbl_cc_library(
+    name = "ValueBoundsOpInterfaceIncGen",
+    strip_include_prefix = "include",
+    tbl_outs = [
+        (
+            ["-gen-op-interface-decls"],
+            "include/mlir/Interfaces/ValueBoundsOpInterface.h.inc",
+        ),
+        (
+            ["-gen-op-interface-defs"],
+            "include/mlir/Interfaces/ValueBoundsOpInterface.cpp.inc",
+        ),
+    ],
+    tblgen = ":mlir-tblgen",
+    td_file = "include/mlir/Interfaces/ValueBoundsOpInterface.td",
+    deps = [
+        ":ValueBoundsOpInterfaceTdFiles",
+    ],
+)
+
+cc_library(
+    name = "ValueBoundsOpInterface",
+    srcs = ["lib/Interfaces/ValueBoundsOpInterface.cpp"],
+    hdrs = ["include/mlir/Interfaces/ValueBoundsOpInterface.h"],
+    includes = ["include"],
+    deps = [
+        ":Analysis",
+        ":DestinationStyleOpInterface",
+        ":IR",
+        ":Support",
+        ":ValueBoundsOpInterfaceIncGen",
+        "//llvm:Support",
+    ],
+)
+
 cc_library(
     name = "TilingInterface",
     srcs = ["lib/Interfaces/TilingInterface.cpp"],
     hdrs = ["include/mlir/Interfaces/TilingInterface.h"],
     includes = ["include"],
     deps = [
         ":DialectUtils",
         ":IR",

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

deps = [
         "//mlir:DialectUtils",
         "//mlir:FuncDialect",
         "//mlir:IR",
         "//mlir:MemRefDialect",
         "//mlir:Pass",
         "//mlir:SCFDialect",
         "//mlir:Support",
         "//mlir:Transforms",
+        "//mlir:ValueBoundsOpInterface",
         "//mlir:VectorDialect",
         "//mlir:VectorUtils",
     ],
 )
 
 cc_library(
     name = "TestDLTI",
     srcs = glob(["lib/Dialect/DLTI/*.cpp"]),