[mlir] Enhance InferShapedTypeOpInterface and move LinalgOps to use them.

[mravishankar]

Summary

or most Linalg operations, the shape of the output can be computed
from the shape of the inputs. A new interface is added that allows
operations to define the shape of the output in terms of shapes of its
operands. A new canonicalization pattern is added to canonicalize dim
ops that query the shape of the result of operations that implement
this interface. This replaces the op -> dim canonicalization
patterns inserted in a piecewise manner.

[mravishankar]

Updating the patch to not use tensor.from_elements. Instead change
to interface to return the Values directly for dims of all results.

[benvanik]

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

1023

can you try using createOrFold here and elsewhere when the dims are created? in theory that'll immediately fold if the dim+source can resolve, preventing the insertion of the dim op entirely

[jpienaar]

It's fine to have these separate for review, but please merge to avoid updating the interface in two consecutive changes where one undos part of the previous

[mravishankar]

Updating by folding in patch D97532.

[mravishankar]

I folded things into a single patch. (Was planning to do that, sorry wasnt explicit before).

@benvanik I saw your comment of using createOrFold (it isnt there on this patch anymore), that made me think we should have SmallVectorImpl<SmallVector<OpFoldResult>> & that might just remove some unnecessary IR being created.

[stellaraccident]

Thanks, this lgtm and it appears that review comments have been addressed.

[jpienaar]

mlir/include/mlir/Interfaces/InferTypeOpInterface.td

115

Every shapedtype only has one shape, this results the value computed for the shape of the shapedtype. Not sure why vector of vectors are needed.

[mravishankar]

Rebase and address concerns from offline conversation

[mravishankar]

Rebase and remove unused methods.

[jpienaar]

mlir/include/mlir/Interfaces/InferTypeOpInterface.td

102

Why this change?

129

Also document that this method is restricted to only ranked cases and cannot be used for unranked contrasted to more general above.

140

The top one uses Value, this one OpFoldResult while the comment says Value. This method is reifying/building the ops to compute the shape. The use of this is in many context, in some of those these could be extracted into pure compute parts. In cases where the type is (for example) the same, would still require to create an Attribute (as the shape is stored in the type, but not as a separate Attribute). So if asked reify the shape for me, if the return is a free standing attribute then the caller would still need to create a constant op.

Lets keep this as Value, as this function is intended for reification. It can be used/but isn't meant as an optimized compute shape by folding & it is also possible to only call reify for dynamic shapes/dims. That way the reify doesn't become mayReify and require consumers to check whether Value or not is returned.

165

Is this change needed? The InferTensorType trait has this method.

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

1568 ↗

(On Diff #329161)

Should this be in StandardOps then? (e.g., making standardops depend on tensor dialect, it does already but not sure if that is intended state there)

1596 ↗

(On Diff #329161)

Should be these be switched? E.g., I can see how one would use a helper method to implement both of these method, and so both interface methods would succeed, but the latter is more efficient if already in that form.

[jpienaar]

Adding Stephan as I'll be OOO

[mravishankar]

mlir/include/mlir/Interfaces/InferTypeOpInterface.td

102

Otherwise every op needs to implement this method. I was trying to make this more opt-in. If a method implements it and an analysis can use it great. But I dont see a need for it to be forced on all operations that implement this interface.

129

Good point. Will do so.

140

I am not sure what is the use of keeping it as Value. I see OpFoldResult is a super set of Value. If the op implementing the interface knows it is a static value, it will return a static value as an attribute. If the method querying this information can use OpFoldResult then its a win-win. A static value is maintained throughout without having to unnecessarily instantiate a Value only for the canonicalizer to make it go away. If the client needs a Value it can always create the Value as needed.

Returning Value and relying on canonicalizers to make them constants is problematic. There are many cases where certain patterns kick in only when shapes are static. For example if you are tiling + vectorizing an operation, then vectorization requires all shapes to be statically known. If the tiling cannot ensure that the shapes are static by construction (and rely on canonicalization to propagate) that information instead it creates a break point. You have to run a pass to do tiling, then run a pass to do vectorization. That leads to other issues cause you need to make sure you are only vectorizing the op you tiled in a previous pass and not other ops (of the same type) that might exist in your program/function that are not meant to be vectorized. THis is just an example where relying on canonlicazations to propagate static information just makes things harder.

Bottom line though is that OpFoldResult is stricly more expressive than Value (its literally a PointerPair<Value, Attribute>). So you can use a Value if thats the case that suits you or an Attribute if you can.

165

It doesnt declare the inferReturnTypeComponents. With the default implementation on this method, by default ODS will not generate the declaration for this method on the op. This gets back to the previous behavior of the defvar.

[herhut]

mlir/include/mlir/Interfaces/InferTypeOpInterface.td

140

I think the question here then is what reify means. If a certain output dimension of an operation has a static shape, then you can just query the type. No need to reify anything, as you have full static knowledge already. If it is a dynamic dimension, you can use this interface to produce IR that makes this dynamic value available to you.

How about changing this interface so that it only produces the shape for a specific dimension and result? That would keep it as a pure reification interface and the logic you want could be in a helper? That is essentially the same argument you gave, just the other way round.

A default implementation could be to fall back to the other interface and insert an extract.

mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp

202

Why the orFold? What does the ConstantIndexOp fold into?

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

1610 ↗

(On Diff #329161)

Is this a good canonicalization in general? In essence, it replicates the computation of the shape in IR instead of falling back on the memref descriptor. That is a good thing to do if you at runtime do not have descriptors. It is less good if you have, as you replicate the computation.

We have the same concept with shape_of and not dim. So generally, shape_of(a) can be replaced with the result of the other shape interface and that is useful in some cases for static analysis purposes and to discharge constraints. However, we do not have this as a canonicalization pattern as it also might duplicate computation. So we have it in a separate pass.

[mravishankar]

Moving the interface into Linalg for now till can converge on using InferShapedTypeOpInterface

[mravishankar]

@herhut Had some offline discussions with Jacques. For the time being moving the interface into Linalg. We can unify them once we iron out all the API of the interface. I am happy to pick that up based on what is finally converged upon.

W.R.T your comment about using OpFoldResult, I think it comes down to what you are using the interface for. To me it just felt more uniform to get the shape of all the dimensions. There is no need to check on the analysis side to check if the shape is dynamic and then go query the interface. You query the interface and it gives you the entire shape (static and dynamic) and then its a matter of just extracting the Value (or int64_t if you can use that) on the client side. Much cleaner usage AFAICS.

W.R.T. having an interface that allows you to get a single dim of a single result, that was one of the things done in an earlier version of the patch. But Jacques pointed out that it is not general enough. Some ops might not be able to compute a single dim without doing all the work to compute the entire dim. In which case if you do need multiple dims of the same result, then you have redundant (potentially heavy) computation that might not CSE. That makes sense to me. There is some API issues to work out here. Landing this in Linalg for now.

[mravishankar]

Fix build file.

[mehdi_amini]

For the time being moving the interface into Linalg. We can unify them once we iron out all the API of the interface. I am happy to pick that up based on what is finally converged upon.

What's the plan to get there? It seems like moving this to linalg is kind of working around the lack of consensus on the API, in practice it seems like it can split the implementation into "Linalg-specific" things and the general case, which will lead to implementations of analysis that are linalg specific. While there are many reasons to have Linalg-specific analysis/transformations, I'm not sure this is such a case (which seems acknowledge by your sentence I'm quoting here).

[mravishankar]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

1610 ↗

(On Diff #329161)

I see what you are saying here. I dont want to start a discussion about whether memref descriptor is a good thing or not, but I dont think it is a good idea for IR canonicalizations to be indexed on how memref descriptors work. If the memref descriptors are replicating computation, then that seems to be an issue with memref descriptor.

Here this pattern is really important. When you replace an operation that has dynamic shapes with another operation, the shape of the replacement obviously needs to match the shape of this op. If you use a the result SSA value, then you cannot remove all uses of the replaced op (and you need to move the insertion point after the operation being replaced to make sure you are no violating use-def chains). Instead resolving the dims in terms of the inputs avoids all of this. This is the reason why I started doing this work. So it is a good canonicalization to have :) . In any case, this is now moved into Linalg dialect only, so it should not affect your use case presently.

[mravishankar]

In D97887#2618111, @mehdi_amini wrote:

For the time being moving the interface into Linalg. We can unify them once we iron out all the API of the interface. I am happy to pick that up based on what is finally converged upon.

What's the plan to get there? It seems like moving this to linalg is kind of working around the lack of consensus on the API, in practice it seems like it can split the implementation into "Linalg-specific" things and the general case, which will lead to implementations of analysis that are linalg specific. While there are many reasons to have Linalg-specific analysis/transformations, I'm not sure this is such a case (which seems acknowledge by your sentence I'm quoting here).

Yeah, I'd like to as well. The hope is that anything in this interface can be made available to the InferShapedTypeOpInterface by just redirecting the call. I think there are larger disagreements about how the shape inference work should evolve. I am not aware of all the opinions here and cannot comment on those.

[mehdi_amini]

I think there are larger disagreements about how the shape inference work should evolve. I am not aware of all the opinions here and cannot comment on those.

I'm not plugged into the details here, but if there is a disagreement it should be addressed. I reiterate what I wrote above: it seems like moving this to Linalg is kind of working around the lack of consensus, do I miss something here?

[mravishankar]

In D97887#2618203, @mehdi_amini wrote:

I think there are larger disagreements about how the shape inference work should evolve. I am not aware of all the opinions here and cannot comment on those.

I'm not plugged into the details here, but if there is a disagreement it should be addressed. I reiterate what I wrote above: it seems like moving this to Linalg is kind of working around the lack of consensus, do I miss something here?

Any suggestions of how to proceed. I am honestly not the person who can drive consensus here cause I do not have the larger context in which the upstream interface is designed. Based on offline discussion this was decided as an intermediate step till the pieces can be merged. This is going to be a simple interface with very fixed functionality. I'll leave it to @nicolasvasilache and @stellaraccident to decide if this is worth landing here or not. I need this functionality so if this is a no-go, I will have to find less diserable WAR solutions (either in Linalg or in IREE).

[stellaraccident]

In D97887#2618411, @mravishankar wrote:

In D97887#2618203, @mehdi_amini wrote:

I think there are larger disagreements about how the shape inference work should evolve. I am not aware of all the opinions here and cannot comment on those.

I'm not plugged into the details here, but if there is a disagreement it should be addressed. I reiterate what I wrote above: it seems like moving this to Linalg is kind of working around the lack of consensus, do I miss something here?

Any suggestions of how to proceed. I am honestly not the person who can drive consensus here cause I do not have the larger context in which the upstream interface is designed. Based on offline discussion this was decided as an intermediate step till the pieces can be merged. This is going to be a simple interface with very fixed functionality. I'll leave it to @nicolasvasilache and @stellaraccident to decide if this is worth landing here or not. I need this functionality so if this is a no-go, I will have to find less diserable WAR solutions (either in Linalg or in IREE).

I can review the thread and try to catch up later, but if there is a conflicting opinion on the shape side, we probably need to resolve it with jpienaar, and i think he is out until next week.

As a side note, this is the cost of IREE being locked to Google's at head development cycle: there is no ability to land a temporary fix past upstream head to make progress. Circt, for example, maintains a project specific LLVM branch for this reason. I know more than anyone that there are reasons for this, but I feel like it is kind of hard for us to hold those reasons and feel urgency about a few days to converge -- when that is a byproduct of Google, not LLVM development process.

Fwiw, the tensorflow/xla side has at times been guilty of this false urgency too, and I would like to see us all be forced to improve in this area.

[mravishankar]

In D97887#2618450, @stellaraccident wrote:

In D97887#2618411, @mravishankar wrote:

In D97887#2618203, @mehdi_amini wrote:

I think there are larger disagreements about how the shape inference work should evolve. I am not aware of all the opinions here and cannot comment on those.

I'm not plugged into the details here, but if there is a disagreement it should be addressed. I reiterate what I wrote above: it seems like moving this to Linalg is kind of working around the lack of consensus, do I miss something here?

Any suggestions of how to proceed. I am honestly not the person who can drive consensus here cause I do not have the larger context in which the upstream interface is designed. Based on offline discussion this was decided as an intermediate step till the pieces can be merged. This is going to be a simple interface with very fixed functionality. I'll leave it to @nicolasvasilache and @stellaraccident to decide if this is worth landing here or not. I need this functionality so if this is a no-go, I will have to find less diserable WAR solutions (either in Linalg or in IREE).

I can review the thread and try to catch up later, but if there is a conflicting opinion on the shape side, we probably need to resolve it with jpienaar, and i think he is out until next week.

Summary of some issues that are unresolved.

1. The current reify... methods need to populate a SmallVector<Value> with as many Values as the number of results of the operation. So each Value is potentially a tensor< rank x index>. This needs to be created by tensor.from_elements in the callee and then the caller needs to tensor.extract the values it needs. Canonicalize and DCE/CSE will eliminate redundant ops. This makes sense for ops where the shape of the result is computed as a single value already, but not for ops, where each dim of each result is computed independently. So the downside here is that a) Everywhere you need to use that interface you need to depend on tensor dialect, which seems unnecessary b) THere is unnecessary cost of adding IR just to remove these instructions in canonicalizations. This is purely an artifact of the interface
2. To account for the above, the next approach was to add a second interface method where the ops implementing the interface need to populate SmallVector<SmallVector<OpFoldResult>>. So there as many vectors as the number of results and each vector is of size rank of the result. The use of OpFoldResult allows returning both static values and dynamic values. There was push back against using OpFoldResult instead of Value. From experience the use of OpFoldResult is really important. Many times (almost always) these shapes returned are used as shapes of results for a replacement operation. Making the interface return Values means that the created operation is now dynamically shaped (even if the original value is statically shaped). To fix up all uses you need to add tensor.cast operations and then rely on canonicalizers to make the result statically shaped again, and remove all the tensor casts. This is again unnecessary addition/deletion of instructions created due to the interface (assuming that the canonicalizers do everything that is needed which is not a given)

I dont want to interfere with what the existing interface is doing. The intermediate interface added to Linalg can always be used to provide the implementation of everything InferShapedTypeOpInterface needs. So I dont see any conflict here.

[stellaraccident]

I feel like I can accept that there is a place for a more specialized shape query interface specifically for linalg, given that it has a more constrained representation, allowing for a narrower interface, compared to the global InferShapeTypeOpInterface.

I'd be inclined to say that the original proposal to add more constrained per-dimension query mechanisms to the global InferShapeTypeOpInterface was not the right direction, since the new methods would only apply to specific ops and circumstances versus being a general purpose shape query mechanism. If we think we need something more specific for a defined subset of cases, then it is the right thing to not try to generalize the global interface in one step. Implementing locally for the case where it is deemed useful seems like the right incremental path forward.

I can take a more detailed final pass if that approach seems reasonable. It looks like Jacques delegated to herhut while he is out: herhut should get a chance to respond when he comes online for his day. It then looks like Mehdi was primarily objecting on non-consensus grounds. If herhut is not opposed to the approach, then neither am I.

mlir/include/mlir/Dialect/Linalg/IR/LinalgInferShapeInterface.td

19 ↗

(On Diff #329790)

If I'm doing a fly-by and I see this interface, the first thing I want to know is "how is this different from the global InferShapeTypeOpInterface?

As it turns out, I do know this area and can glean the difference, and I wonder if some name clarification (or comments) can help.

The global interface is generic over any shape transfer function, including, I believe, ranked, unranked and dynamic. This interface is providing a narrow facility to query individual dimensions of the results. I can see the arguments about why you would want this more direct access for various things, given that the ops here are constrained such that it is always valid to service such a query, whereas the global InferShapeTypeOpInterface is higher level (you can implement InferShapeTypeOpInterface in terms of InferShapeOp here).

If we go forward with this, it makes more sense to me for this interface to be something like InferResultDimsInterface and to make the method a bit cleaner (comment below).

41 ↗

(On Diff #329790)

In what situations can this fail?

42 ↗

(On Diff #329790)

getResultShapeDims if renaming.

44 ↗

(On Diff #329790)

I don't love this calling convention. Up thread, Jacques had asked (on the original) version, whether we wanted to return just dims for a certain result index or result+dim. In general, we've biased towards not economizing at the individual dim level (and as you have it, it can return an attribute for static dims, eliminating most practical redundancies). Any chance we could at least do this by result index and avoid the double vector? I don't feel super strongly about this - more of a smell.

[mravishankar]

Document the interface scope better.

mlir/include/mlir/Dialect/Linalg/IR/LinalgInferShapeInterface.td

19 ↗

(On Diff #329790)

Yeah this makes sense. That was feedback too, that I was maybe looking at optimizing for a specific use case. So the intent is to have a Linalg specific interface that is for this optimized used case, that can also tie into the larger InferShapedTypeOpInterface. I added an explicit note stating this here (kept the name same for now, will change it based on other feedback).

41 ↗

(On Diff #329790)

The failure is for two uses case. One is based on how the next iteration of this might look like (it is described below), but another reason is that this interface (I think) is useful when the shape of the result of the operation is expressible using the shape of the operands. The failure is to indicate that this is not possible for the op. Today it is not true for any Linalg operation (and shouldnt be true for any operation that is tensor based AFAICS). Still having an escape hatch here will help exit gracefully when failure does happen (instead of asserts that might not trigger based on build)

42 ↗

(On Diff #329790)

(see below)

44 ↗

(On Diff #329790)

The thought here is that each op might have different constraints w.r.t how the shapes need to be computed. It might be not possible (or at least not optimal) for an op to compute the shape of every result dim individually. So there is a natural heirarchy of methods (each of which returns failure by default)

a) getResultShapes: Return the shape of all dims for all results
b) getResultShape (singular): Return the shape of all dims of a particular result
c) getResultDim : Return the shape of a particular dim of a particular result.

An op can override any of these methods. The default implementation of (a) uses (b), and the default implementation of (b) uses (c).

From the perspective of the analysis/transformation the hierarchy goes in reverse. So we need a set of utility functions (Static Interface methods) that use the above interface as following
a) computeResultShapes(InferShapeOp op) : Returns the the result of op.getResultShape
b) computeResultShape(InferShapeOp op, unsigned index): Returns the result at position index of the op. If op implements getResultShape use that, otherwise use the result of computeResultShape and extract the needed values.
c) computeResultDim(InferShapeOp op, unsigned index, unsigned dim) : return the shape of a particular result, dim. If the op implements getResultDim use that, if not use computeResultShape and extract the needed values.
(i.e. (c) uses (b) as fallback and (b) uses (a) as fallback, i.e reverse of the above get* methods).

All this is just an idea, and planning to implement only on need basis. For now. Just getResultShape will suffice to see if we need this whole layering, either for optimization or correctness. All this was to justify why I went with the current calling convention. It is the least opinionated.

[mravishankar]

mlir/include/mlir/Dialect/Linalg/IR/LinalgInferShapeInterface.td

44 ↗

(On Diff #329790)

Typo:

a) computeResultShapes(InferShapeOp op) : Returns the the result of op.getResultShapes

[mravishankar]

mlir/include/mlir/Dialect/Linalg/IR/LinalgInferShapeInterface.td

44 ↗

(On Diff #329790)

Another typo :

All this is just an idea, and planning to implement only on need basis. For now. Just getResultShapes will suffice to see if we need this whole layering, either for optimization or correctness. All this was to justify why I went with the current calling convention. It is the least opinionated.

[stellaraccident]

This seems fine to me, and I believe that Jacques is supportive of a local/specialized thing here. However, since he has open comments, I would prefer that we get his lgtm on Monday when he is back prior to proceeding.

[nicolasvasilache]

I have not yet dug into the guts of the discussion and the implementation details, I'll just comment on 2 things:

There was push back against using OpFoldResult instead of Value. From experience the use of OpFoldResult is really important.

Where it makes sense, it is crucial to move towards a ValueOrAttr abstraction rather than stay stuck in Value land.
Maybe it does not (yet?) make sense for the Shape world and that is perfectly fine.
However this is where Linalg is: see in particular the getMixed methods in ViewLikeInterface.
The quality of life improvements brought about by that particular refactoring were enormous.

I've also seen similar discussions kick in in a vector context: one want to use ValueOrAttr aggressively to get a static constant that can be used to build a new vector shape.
Any API in the middle that does not use ValueOrAttr when it could becomes a (multiplicative) efficiency loss.

I reiterate what I wrote above: it seems like moving this to Linalg is kind of working around the lack of consensus, do I miss something here?

I don't see a particular need for consensus between the needs of Linalg and the solution provided by the Shape dialect at this point in spacetime.
What makes sense for Linalg today may or may not make sense for Shape tomorrow.

Stepping back from Linalg, I think we should generally stop wanting to solve all problems at the same time in the same abstraction: this is a recurrent source of a lot of wasted effort.
The ultimate goal is of course convergence but it must happen progressively through iteration.
Progress is not a step function: 0: the wheel 1: flying cars.

The risk is of course time spent in refactorings.
Yes it can be expensive; yet it is absolutely dwarfed by the cost of inertia and missed opportunities.

[herhut]

In D97887#2618054, @mravishankar wrote:

W.R.T your comment about using OpFoldResult, I think it comes down to what you are using the interface for. To me it just felt more uniform to get the shape of all the dimensions. There is no need to check on the analysis side to check if the shape is dynamic and then go query the interface. You query the interface and it gives you the entire shape (static and dynamic) and then its a matter of just extracting the Value (or int64_t if you can use that) on the client side. Much cleaner usage AFAICS.

I would have hidden this behind a helper that wraps the API, so that the reify function always reifies but only gets called if the shape is dynamic. However...

W.R.T. having an interface that allows you to get a single dim of a single result, that was one of the things done in an earlier version of the patch. But Jacques pointed out that it is not general enough. Some ops might not be able to compute a single dim without doing all the work to compute the entire dim. In which case if you do need multiple dims of the same result, then you have redundant (potentially heavy) computation that might not CSE. That makes sense to me. There is some API issues to work out here. Landing this in Linalg for now.

... I did not think about this. I had simply assumed that reifying multiple times ultimately is free. If it is not, as you suggest, then reifying single dimensions is also not possible. And, consequently, the API can not be wrapped as nicely. In particular, it would potentially materialize shape computations that are not needed and cannot be CSE'd. The same argument again.

So if we assume that shape computations are not CSE-able, then something like OpFoldResult makes a lot of sense. I would give it a new name, though, but that is pure stylistics.

[stellaraccident]

In D97887#2621515, @nicolasvasilache wrote:

I reiterate what I wrote above: it seems like moving this to Linalg is kind of working around the lack of consensus, do I miss something here?

I don't see a particular need for consensus between the needs of Linalg and the solution provided by the Shape dialect at this point in spacetime.
What makes sense for Linalg today may or may not make sense for Shape tomorrow.

Fwiw, this isn't about shape dialect+linalg, the original patch was extending a global interface for shape inference which has claims that seemed overlapping. I think that the point has been resolved that these are different (in that the linalg version is a specialization that we don't necessarily want on the global one). My pushback on landing until Jacques or Stephan weighed in was that they had taken the time to do detailed reviews of the original and we should provide time for the conversation to close (and I know that Jacques is on vacation because he pinged me privately prior to heading out).

[mehdi_amini]

In D97887#2621515, @nicolasvasilache wrote:

I reiterate what I wrote above: it seems like moving this to Linalg is kind of working around the lack of consensus, do I miss something here?

I don't see a particular need for consensus between the needs of Linalg and the solution provided by the Shape dialect at this point in spacetime.
What makes sense for Linalg today may or may not make sense for Shape tomorrow.

If we're aiming upstream to build interfaces and general solution to manipulate shape inference (reification or not), then I expect that we dogfood. Yes there is a cost, but that has to be the cost of developing linalg upstream.

I always mention to people that we only have the concept of OpInterface in MLIR in the first *because* you had a need in linalg and MLIR didn't have a solution for this.

Stepping back from Linalg, I think we should generally stop wanting to solve all problems at the same time in the same abstraction: this is a recurrent source of a lot of wasted effort.
The ultimate goal is of course convergence but it must happen progressively through iteration.
Progress is not a step function: 0: the wheel 1: flying cars.

Yes: but progress is not a random walk either. You can think about the end goal and have an idea about whether your next step is getting you closer to the flying car ;)
And as a community that joins multiple teams / group with various deployment/integration story, sharing this understanding is a key prerequisite to iterate in the project.

Avoiding to build the vision about the direction, and skip these discussions for the sake of your "island's velocity" does not sit well with me: for example your experience with Value / OpFoldResult is very important to be discussed and understood by others who don't work daily on linalg, this can impact the rest of the project. The points Mahesh is making in this thread make sense to me, and the most recent exchange with Stephan for example shows that discussing this is worthwhile in evolving our shared understanding of the tradeoffs.

The risk is of course time spent in refactorings.
Yes it can be expensive; yet it is absolutely dwarfed by the cost of inertia and missed opportunities.

I don't think it is that black or white, and I think that caution in the consensus when working upstream is critical to the project: I'm personally wary of building a "tech island" around linalg, with a lot of things that would not interface with the rest of the ecosystem.

[jpienaar]

In D97887#2621515, @nicolasvasilache wrote:

I have not yet dug into the guts of the discussion and the implementation details, I'll just comment on 2 things:

There was push back against using OpFoldResult instead of Value. From experience the use of OpFoldResult is really important.

Where it makes sense, it is crucial to move towards a ValueOrAttr abstraction rather than stay stuck in Value land.
Maybe it does not (yet?) make sense for the Shape world and that is perfectly fine.

I find it funny how folks (excl Stephan) miss something very simple here: the method that was being changed is called reify. It reifies the shape computation, it doesn't potential reify. So Attributes doesn't makes sense. The caller expects a Value corresponding to the reified shape computation. This would be like having a build method and the build may or may not build (even createOrFold returns a Value rather than an attribute), or materializeConstant which decides to not materialize. It means the caller would need to check which state is made and then create a constant op or not, so this pushes work on to the caller if we were to relax it.

Attributes aren't free here either: if you wanted to avoid overhead during inference an int or an arrayref<int> or ShapedComponentType is better as for intermediate shapes during inference or when using existing shapes from a ShapedType, you don't want to create an attribute. That is uniqueing it in the context without need, it never needs to be a attribute (it isn't one in the final result type contrary to folding where it may be part of the constant op finally created, so there it is only potentially wasteful to make an attribute, here it would always be). But again that is a different method than reify.

In D97887#2621926, @herhut wrote:

In D97887#2618054, @mravishankar wrote:

W.R.T. having an interface that allows you to get a single dim of a single result, that was one of the things done in an earlier version of the patch. But Jacques pointed out that it is not general enough. Some ops might not be able to compute a single dim without doing all the work to compute the entire dim. In which case if you do need multiple dims of the same result, then you have redundant (potentially heavy) computation that might not CSE. That makes sense to me. There is some API issues to work out here. Landing this in Linalg for now.

... I did not think about this. I had simply assumed that reifying multiple times ultimately is free. If it is not, as you suggest, then reifying single dimensions is also not possible. And, consequently, the API can not be wrapped as nicely. In particular, it would potentially materialize shape computations that are not needed and cannot be CSE'd. The same argument again.

Yes so for cases where computing a single dim requires computing all the dims of the shape, it would seem more fragile (and potentially more expensive) to rely on CSE than having 0-1 tensor.concat (per result) or shape.from_extents and relying on that & get_element_at_dim (forgot the name) to be able to compose - the objection was that one now has (potentially) a concat & multiple get element calls which get inserted and removed later vs forwarding directly from the start (but again, thats an easy pattern that createOrFold could even do, and then it is just deleting the potentially optional concat means at most 1 extra op here). And good point about wrapping, I didn't consider that end.

For LinAlg, is this also true? E.g., is the output shape indices interrelated? If all of them are indeed calculatable per dim, then that changes things here given interface is LinAlg specific now.

This is fine for local iteration in LinAlg, it can't be used by shape inference pass and I don't think the tradeoff of at most one op per result that needs to be inserted and deleted is worth it really (given how often this would be called - one reify call per bufferization?).

mlir/include/mlir/Dialect/Linalg/IR/LinalgInferShapeInterface.td

41 ↗

(On Diff #329790)

So failure here could mean either input shapes are invalid or can't get the shape as attributes and values? InferShapedTypeOpInterface is used along with InferTypeOpInterface (mostly, not always, so perhaps I should expand the documentation there to make it clearer) so the verifier catches all invalid shapes so that the reify methods return is less ambiguous (e.g., failure cannot imply invalid input shapes as a verified op can't have it). Are the shapes for LinAlg ops verified elsewhere?

44 ↗

(On Diff #329790)

I agree with Stella here, but that is open for you two to discuss too.

[mravishankar]

In D97887#2632735, @jpienaar wrote:

In D97887#2621515, @nicolasvasilache wrote:

I have not yet dug into the guts of the discussion and the implementation details, I'll just comment on 2 things:

There was push back against using OpFoldResult instead of Value. From experience the use of OpFoldResult is really important.

Where it makes sense, it is crucial to move towards a ValueOrAttr abstraction rather than stay stuck in Value land.
Maybe it does not (yet?) make sense for the Shape world and that is perfectly fine.

I find it funny how folks (excl Stephan) miss something very simple here: the method that was being changed is called reify. It reifies the shape computation, it doesn't potential reify. So Attributes doesn't makes sense. The caller expects a Value corresponding to the reified shape computation. This would be like having a build method and the build may or may not build (even createOrFold returns a Value rather than an attribute), or materializeConstant which decides to not materialize. It means the caller would need to check which state is made and then create a constant op or not, so this pushes work on to the caller if we were to relax it.

Attributes aren't free here either: if you wanted to avoid overhead during inference an int or an arrayref<int> or ShapedComponentType is better as for intermediate shapes during inference or when using existing shapes from a ShapedType, you don't want to create an attribute. That is uniqueing it in the context without need, it never needs to be a attribute (it isn't one in the final result type contrary to folding where it may be part of the constant op finally created, so there it is only potentially wasteful to make an attribute, here it would always be). But again that is a different method than reify.

I think that is a good reason to make this Linalg only. In Linalg based codegen not having static shapes has real issues. So this interface addresses those issues. There is more work needed here to bridge the gap between what is needed here and what the Shape inference interface needs.

In D97887#2621926, @herhut wrote:

In D97887#2618054, @mravishankar wrote:

W.R.T. having an interface that allows you to get a single dim of a single result, that was one of the things done in an earlier version of the patch. But Jacques pointed out that it is not general enough. Some ops might not be able to compute a single dim without doing all the work to compute the entire dim. In which case if you do need multiple dims of the same result, then you have redundant (potentially heavy) computation that might not CSE. That makes sense to me. There is some API issues to work out here. Landing this in Linalg for now.

... I did not think about this. I had simply assumed that reifying multiple times ultimately is free. If it is not, as you suggest, then reifying single dimensions is also not possible. And, consequently, the API can not be wrapped as nicely. In particular, it would potentially materialize shape computations that are not needed and cannot be CSE'd. The same argument again.

Yes so for cases where computing a single dim requires computing all the dims of the shape, it would seem more fragile (and potentially more expensive) to rely on CSE than having 0-1 tensor.concat (per result) or shape.from_extents and relying on that & get_element_at_dim (forgot the name) to be able to compose - the objection was that one now has (potentially) a concat & multiple get element calls which get inserted and removed later vs forwarding directly from the start (but again, thats an easy pattern that createOrFold could even do, and then it is just deleting the potentially optional concat means at most 1 extra op here). And good point about wrapping, I didn't consider that end.

For LinAlg, is this also true? E.g., is the output shape indices interrelated? If all of them are indeed calculatable per dim, then that changes things here given interface is LinAlg specific now.

It is true for some ops, not true for others. The major case is for structured ops. You could compute the shape indices for each element. There is some efficiencies that can be exploited for Structured ops if you compute the shape of all results at once. For other ops like PadTensorOp and InitTensorOp, etc, they difference indices are computed completely independently. So the default here is to compute everything at once. Depending on use cases we can make things a bit more fine-grained.

This is fine for local iteration in LinAlg, it can't be used by shape inference pass and I don't think the tradeoff of at most one op per result that needs to be inserted and deleted is worth it really (given how often this would be called - one reify call per bufferization?).

[mravishankar]

Going back to the implementation which updated the interface of
InferShapedTypeInterface but dropping the use of OpFoldResult.

[mravishankar]

mlir/include/mlir/Interfaces/InferTypeOpInterface.td

129

I can update the doc, but it seems like it could be implemented by unranked to (which I am assuming is <*xf32>? Sorry, not sure if that is unranked, but happy to update the doc.

140

Ok, I changed it to use Value

mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp

202

Probably nothing. Hoping that it would return a previous constant operation of the same value. Doesnt do. Dont see a harm in leaving it that way.

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

1023

Using the "mixed mode" avoids the issue now.

[stellaraccident]

I'm having a bit of trouble following the back and forth on this patch (linalg specific vs touching the global InferShapeType interface). I left some specific comments. I don't have an opinion (and defer to those who do) on the "reify" and OpFoldResult vs Value points.

I feel like the description should be updated to more clearly describe where this is landing.

mlir/include/mlir/Interfaces/InferTypeOpInterface.td

129

I can update the doc, but it seems like it could be implemented by unranked to (which I am assuming is <*xf32>? Sorry, not sure if that is unranked, but happy to update the doc.

I think that since this is returning a vector of dims, it is kind of implied that it only applies to ranked cases. Typically if returning a "shape" as per the above/pre-existing methods, an unranked result type would/could return a tensor<?xindex> (or equiv) as its "shape" (and have some kind of whole tensor operations that calculate it in a way that does not pre-suppose knowledge of a rank). With the way this is defined, you lose that generality, and therefore, it is restricted to ranked only.

mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp

202

I would make it just create to avoid confusion of people reading and wondering if they are missing something (I paused and tried to figure out what this could possibly be doing for a constant - better if that cognitive overhead was removed).

(also for consistency with the rest of the patch)

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

680

Should this pattern have a dedicated unit test (vs being tested incidentally via a linalg op)?

697

I feel that this could use a comment describing when this cascade (first checking reifyReturnTypeShapes -> reifyReturnTypeShapesPerResultDim) falls through to the second case.

713

Nit: Maybe call it reifiedReturnDims? One of the things that is confusing me about this patch is that historically, when we say "shape" we refer to a single SSA value that is either a !shape.shape type or a tensor<?xindex>, whereas these per dim variants are inconsistently referring to static lists of dims as "shapes".

Also: You define a differently typed temporary in the block above with the same name. Please don't shadow names like this.

719

Is there a case where we can get here and the result does not have a rank?

mlir/lib/Dialect/StandardOps/CMakeLists.txt

18 ↗

(On Diff #331758)

I think you can revert the changes to this file?

mlir/lib/Interfaces/InferTypeOpInterface.cpp

14–16

NFC this cleanup separately?

[jpienaar]

LG to me interface wise modulo some ergonomics and questions others have there. Thanks!

mlir/include/mlir/Interfaces/InferTypeOpInterface.h

19 ↗

(On Diff #331758)

Is this needed?

mlir/include/mlir/Interfaces/InferTypeOpInterface.td

109

Just to check, so both may be overriden but at least one has to?

The ideal here would be that one can define one or both, but the caller can call either - e.g., callers need not know which one is implemented and folks can override either, now I don't know if we could use some C++ (potentially) template magic to verify this vs having a mutually recursive call in the non-overridden case. Do you have an idea?

[mravishankar]

Addressing comments.

[mravishankar]

mlir/include/mlir/Interfaces/InferTypeOpInterface.h

19 ↗

(On Diff #331758)

THanks. Using just Value this isnt needed anymore.

mlir/include/mlir/Interfaces/InferTypeOpInterface.td

109

I think its better only one is overridden (changed the description to say so). I dont know of a C++ way of enforcing this.

Based on the comments below, it looks like the existing method is strictly more general than the one added in this patch. So the default implementation of the existing method could be to call the method added here. THe part that is unclear to me is how to implement this fall back without adding a dependence on the tensor dialect. On a previous patch, adding dependence of an interface on a dialect was flagged as an issue, particularly because it can cause circular dependencies in the build. If an op from tensor dialect needs to use this interface then it will depend on the library here, and this will in-turn depend on the tensor dialect. I am happy to implement the fall backs as needed.

129

Thanks for the explanation. This makes sense. I updated the description to state that this only supports ranked case.

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

680

Added a new op to test dialect to test the canonicalization pattern outside of Linalg.

697

This is what the comment for the function itself is describing. I moved the comment here where its probably more relevant.

719

Probably not, but this is more of a sanity check. Ideally, would like to verify that this doesnt happen for an op that does implement this method, but the verification will have to introduce operations into the IR, so dont see a good way to verify. For now this canonicalization will just fail gracefully and not just crash.

mlir/lib/Dialect/StandardOps/CMakeLists.txt

18 ↗

(On Diff #331758)

Thanks!

mlir/lib/Interfaces/InferTypeOpInterface.cpp

14–16

With use of Value I dont need this anymore.

[mravishankar]

Rebase

mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td

InterfaceMethod<
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
         if (inputIdx >= getNumInputs()) return {};
         return getOperandDimPositionInLoopsToShapeMap(inputIdx, dim);
       }]
     >,
     InterfaceMethod<
       /*desc=*/[{
-        Return the position in the results of the affine map computed
-        by getLoopsToShapesMap() that represents the shape of the
-        result value at a dimension.
+        Return the range of position in the result of the affine map
+        computed by getLoopsToShapesMap() which correspond to the
+        AffineExprs used to access the outputs of the operation.
       }],
-      /*retTy=*/"Optional<unsigned>",
-      /*methodName=*/"getResultValueDimPositionInLoopsToShapeMap",
-      /*args=*/(ins "unsigned":$resultIdx, "unsigned":$dim),
+      /*retTy=*/"std::pair<unsigned, unsigned>",
+      /*methodName=*/"getResultsPositionInLoopsToShapeMap",
+      /*args=*/(ins),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
-        if (resultIdx >= getNumOutputs()) return {};
-        return getOperandDimPositionInLoopsToShapeMap(
-            getNumInputs() + resultIdx, dim);
+        return
+          {*getOperandDimPositionInLoopsToShapeMap(getNumInputs(), 0),
+           (*getOperandDimPositionInLoopsToShapeMap
+                 (getNumInputs() + getNumOutputs() - 1,
+                  getOutputShapedType(getNumOutputs()-1).getRank() - 1)) + 1};
       }]
     >,
     InterfaceMethod<
       /*desc=*/[{
         Like `getShape`, but only returns statically-known information, without
         generating any new IR. For each shape dimension, returns >=0 if that
         dimension is statically known, or ShapeType::kDynamicSize otherwise.
       }],
Operation::operand_range getAssumedNonShapedOperands() {
         assert((t.isSignlessIntOrIndexOrFloat() || t.template isa<VectorType>())
                &&"expected scalar or vector type");
       }
       return res;
     }
 
     /// Returns the value that expresses the shape of the output in terms of
     /// shape of the input operands where possible
-    Optional<Value> inferResultDimFromInputShapes
-      (OpBuilder &b, Location loc, unsigned resultIdx, unsigned im);
+    LogicalResult reifyReturnTypeShapesPerResultDim(OpBuilder &b,
+        SmallVectorImpl<SmallVector<Value>> &reifiedReturnShapes);
 
     //========================================================================//
     // Helper functions to mutate the `operand_segment_sizes` attribute.
     // These are useful when cloning and changing operand types.
     //========================================================================//
     void setNumInputs(unsigned num) { setOperandSegmentAt(0, num); }
     void setNumOutputBuffers(unsigned num) { setOperandSegmentAt(1, num); }
 

mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.h

 #include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinDialect.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/IR/Types.h"
 #include "mlir/Interfaces/CopyOpInterface.h"
+#include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "mlir/Interfaces/ViewLikeInterface.h"
 #include "mlir/Support/LLVM.h"
 
 #include "llvm/ADT/STLExtras.h"
 
 namespace mlir {
 namespace linalg {

 SmallVector<AffineExpr, 4> concat(ArrayRef<AffineExpr> a,
                                   ArrayRef<AffineExpr> b);
 
 /// Return the dims that are `iteratorTypeName` loops in the LinalgOp `op`.
 /// Assumes `op` is a LinalgOp.
 void getDimsOfType(Operation *op, StringRef iteratorTypeName,
                    SmallVectorImpl<AffineExpr> &res);
 
-/// For reshape operation, compute the shape of the output based on the result
-/// type and shape of the input.
-SmallVector<Value, 4>
-getReshapeOutputShapeFromInputShape(OpBuilder &b, Location loc, Value src,
-                                    ArrayRef<int64_t> dstStaticShape,
-                                    ArrayRef<AffineMap> reassociation);
-
 namespace detail {
 LogicalResult verifyStructuredOpInterface(Operation *op);
 } // namespace detail
 } // namespace linalg
 } // namespace mlir
 
 namespace mlir {
 namespace linalg {

mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LINALG_OPS
 #define LINALG_OPS
 
 include "mlir/Dialect/Linalg/IR/LinalgBase.td"
 include "mlir/Interfaces/ControlFlowInterfaces.td"
+include "mlir/Interfaces/InferTypeOpInterface.td"
 include "mlir/Interfaces/LoopLikeInterface.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 include "mlir/Interfaces/ViewLikeInterface.td"
 
 // Base class for Linalg dialect ops that do not correspond to library calls.
 class Linalg_Op<string mnemonic, list<OpTrait> traits = []> :
     Op<Linalg_Dialect, mnemonic, traits> {
   // For every linalg op, there needs to be a:
   //   * void print(OpAsmPrinter &p, ${C++ class of Op} op)
   //   * LogicalResult verify(${C++ class of Op} op)
   //   * ParseResult parse${C++ class of Op}(OpAsmParser &parser,
   //                                         OperationState &result)
   // functions.
   let printer = [{ return ::print(p, *this); }];
   let verifier = [{ return ::verify(*this); }];
   let parser = [{ return ::parse$cppClass(parser, result); }];
 }
 
-def Linalg_InitTensorOp : Linalg_Op<"init_tensor", [NoSideEffect]> {
+def Linalg_InitTensorOp : Linalg_Op<"init_tensor",
+    [NoSideEffect,
+     DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
+         ["reifyReturnTypeShapesPerResultDim"]>]> {
   let summary = "operation to define a tensor of particular value";
 
   let description = [{
     `linalg.init_tensor` is an operation that materializes a tensor of
     a given shape. The shape could be dynamic or static.
   }];
 
   let arguments =
let builders = [
     OpBuilder<(ins "ArrayRef<OpFoldResult>":$sizes, "Type":$elementType,
       CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>
   ];
 
   let hasCanonicalizer = 1;
 }
 
 def Linalg_PadTensorOp : Linalg_Op<"pad_tensor",
-    [AttrSizedOperandSegments, NoSideEffect]> {
+    [AttrSizedOperandSegments,
+     DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
+         ["reifyReturnTypeShapesPerResultDim"]>,
+     NoSideEffect]> {
   let summary = "tensor pad operation";
   let description = [{
     `linalg.pad_tensor` is an operation that pads the `source` tensor
     with given `low` and `high` padding config.
 
     The PadTensor operation supports the following arguments:
 
     * source: the "base" tensor on which to pad.
code commonExtraClassDeclaration = [{
     SmallVector<ReassociationExprs, 4> getReassociationExprs() {
       return
         llvm::to_vector<4>(llvm::map_range(reassociation(),
           [](Attribute a) {
             return llvm::to_vector<2>(
               a.cast<AffineMapAttr>().getValue().getResults());
           }));
     }
-    SmallVector<Value, 4> getOutputShape(OpBuilder &b, Location loc) {
-      return getReshapeOutputShapeFromInputShape(
-          b, loc, src(), getResultType().getShape(),
-          getReassociationMaps());
-    }
   }];
   let assemblyFormat = [{
     $src $reassociation attr-dict `:` type($src) `into` type(results)
   }];
 }
 
 def Linalg_ReshapeOp : Linalg_ReshapeLikeOp<"reshape",
     [DeclareOpInterfaceMethods<ViewLikeOpInterface>]>,
def Linalg_ReshapeOp : Linalg_ReshapeLikeOp<"reshape",
   let extraClassDeclaration = commonExtraClassDeclaration # [{
     MemRefType getSrcType() { return src().getType().cast<MemRefType>(); }
     MemRefType getResultType() { return result().getType().cast<MemRefType>(); }
   }];
   let hasFolder = 1;
   let hasCanonicalizer = 1;
 }
 
-def Linalg_TensorReshapeOp : Linalg_ReshapeLikeOp<"tensor_reshape">,
+def Linalg_TensorReshapeOp : Linalg_ReshapeLikeOp<
+    "tensor_reshape",
+    [DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
+      ["reifyReturnTypeShapesPerResultDim"]>]>,
     Arguments<(ins AnyTensor:$src,
                    AffineMapArrayAttr:$reassociation)>,
     Results<(outs AnyTensor:$result)> {
   let summary = "linalg.tensor_reshape produces a new reshaped tensor.";
   let description = [{
     The `linalg.reshape` op produces a new tensor whose sizes are a
     reassociation of the original `src`.
 

mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td

 //===----------------------------------------------------------------------===//
 
 #ifndef LINALG_STRUCTURED_OPS
 #define LINALG_STRUCTURED_OPS
 
 include "mlir/Dialect/Linalg/IR/LinalgBase.td"
 include "mlir/Dialect/Linalg/IR/LinalgInterfaces.td"
 include "mlir/Interfaces/CopyOpInterface.td"
+include "mlir/Interfaces/InferTypeOpInterface.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 
 // Base Tablegen class for Linalg ops.
 // Linalg ops that correspond to library calls operate on ShapedType as their
 // first operands. These may be optionally followed by non-view operands
 // depending on the specific Linalg op.
 class LinalgStructuredBase_Op<string mnemonic, list<OpTrait> props>
   : Op<Linalg_Dialect, mnemonic, !listconcat(props, [
-       LinalgStructuredInterface])> {
+       LinalgStructuredInterface, InferShapedTypeOpInterface])> {
   code structuredOpsBaseDecls = [{
     // Return the number of induction variables in the basic block. This should
     // always be 0 for index-free linalg ops. For IndexedGeneric, this must be
     // equal to numLoops.
     unsigned getNumPayloadInductionVariables() {
       return isa<IndexedGenericOp>(this->getOperation()) ? getNumLoops() : 0;
     }
+
+    LogicalResult reifyReturnTypeShapesPerResultDim(OpBuilder &b,
+        SmallVectorImpl<SmallVector<Value>> &reifiedReturnShapes) {
+      return cast<LinalgOp>(getOperation()).reifyReturnTypeShapesPerResultDim(b,
+          reifiedReturnShapes);
+    }
   }];
 }
 
 class LinalgStructured_Op<string mnemonic, list<OpTrait> props>
   : LinalgStructuredBase_Op<mnemonic,
        !listconcat(props, [
          DeclareOpInterfaceMethods<MemoryEffectsOpInterface>])> {
   code structuredOpsDecls = structuredOpsBaseDecls # [{

mlir/include/mlir/Interfaces/InferTypeOpInterface.td

StaticInterfaceMethod<
       /*retTy=*/"::mlir::LogicalResult",
       /*methodName=*/"inferReturnTypeComponents",
       /*args=*/(ins "::mlir::MLIRContext*":$context,
                     "::mlir::Optional<::mlir::Location>":$location,
                     "::mlir::ValueRange":$operands,
                     "::mlir::DictionaryAttr":$attributes,
                     "::mlir::RegionRange":$regions,
                     "::mlir::SmallVectorImpl<::mlir::ShapedTypeComponents>&":
-                      $inferredReturnShapes)
+                      $inferredReturnShapes),
+      /*methodBody=*/[{}],
+      /*defaultImplementation=*/[{ return ::mlir::failure(); }]
     >,
     InterfaceMethod<
       /*desc=*/[{Reify the shape computation for the operation.
 
-      Insert operations using the given OpBuilder that computes the result
-      shape.
+      Insert operations using the given OpBuilder that computes the
+      result shape. Only one of this method or
+      `reifyReturnTypeShapesPerResultDim` needs to be overriden by the
+      operation.
       }],
       /*retTy=*/"::mlir::LogicalResult",
       /*methodName=*/"reifyReturnTypeShapes",
       /*args=*/(ins "::mlir::OpBuilder&":$builder,
-                    "::mlir::SmallVectorImpl<::mlir::Value>&":$reifiedReturnShapes),
+          "::mlir::SmallVectorImpl<Value> &":$reifiedReturnShapes),
       /*methodBody=*/[{}],
       /*defaultImplementation=*/[{ return ::mlir::failure(); }]
     >,
+    InterfaceMethod<
+      /*desc=*/[{Reify the shape computation for the operation.
+
+      Insert operations using the given OpBuilder that computes the
+      result shape. The `reifiedReturnShapes` is expected to be
+      populated with as many vectors as the number of results of the
+      op (empty if the shape of a result value cannot be computed). If
+      the returned shape for a result is not empty, its size must
+      match the rank of the shaped type returned. Consequently, this
+      interface can only be overridden if the return types are ranked.
+
+      If both this method and `reifyReturnTypeShapes` are overridden
+      by the operation, `reifyReturnTypeShapes` takes precedence. This
+      method is intended to be used when the shape of each result, dim
+      pair can be computed independently. Using this method avoids
+      adding additional instructions to aggregate individual dimension
+      of a result shape into an single `Value` (and consequently
+      avoids the need to extract the value from the shape on the
+      client side).
+      }],
+      /*retTy=*/"::mlir::LogicalResult",
+      /*methodName=*/"reifyReturnTypeShapesPerResultDim",
+      /*args=*/(ins "::mlir::OpBuilder&":$builder,
+          "::mlir::SmallVectorImpl<SmallVector<::mlir::Value>>&"
+          :$reifiedReturnShapes),
+      /*methodBody=*/[{}],
+      /*defaultImplementation=*/[{ return ::mlir::failure(); }]
+    >
   ];
 }
 
 // Convenience class grouping together type and shaped type op interfaces for
 // ops that have tensor return types.
 class InferTensorType<list<string> overridenMethods = []> {
   list<OpTrait> traits = [
     // Op implements infer type op interface.
     InferTypeOpInterface,
     // The op will have methods implementing the ShapedType type inference
     // interface.
     DeclareOpInterfaceMethods<InferShapedTypeOpInterface, overridenMethods>,
     // The op produces tensors and will use the ShapedType type infer interface
     // along with knowledge that it is producing Tensors to infer the type.
     NativeOpTrait<"InferTensorType">
   ];
 }
-defvar InferTensorTypeWithReify = InferTensorType<["reifyReturnTypeShapes"]>;
+defvar InferTensorTypeWithReify = InferTensorType<[
+    "inferReturnTypeComponents", "reifyReturnTypeShapes"]>;
 
 #endif // MLIR_INFERTYPEOPINTERFACE

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

 add_mlir_dialect_library(MLIRLinalg
   LinalgInterfaces.cpp
   LinalgOps.cpp
   LinalgTypes.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Linalg
 
   DEPENDS
   MLIRLinalgInterfacesIncGen
   MLIRLinalgOpsIncGen
   MLIRLinalgStructuredOpsIncGen
 
   LINK_LIBS PUBLIC
   MLIRAffine
   MLIRDialectUtils
+  MLIRInferTypeOpInterface
   MLIRIR
   MLIRParser
   MLIRSideEffectInterfaces
   MLIRViewLikeInterface
   MLIRStandard
   MLIRMemRef
   MLIRTensor
   )

mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp

 }
 
 SmallVector<Value, 4> LinalgOp::createFlatListOfOperandDims(OpBuilder &b,
                                                             Location loc) {
   SmallVector<Value, 4> res;
   for (Value v : getShapedOperands()) {
     ShapedType t = v.getType().template cast<ShapedType>();
     for (unsigned i = 0, e = t.getRank(); i < e; ++i)
-      res.push_back(b.create<memref::DimOp>(loc, v, i));
+      res.push_back(b.createOrFold<memref::DimOp>(loc, v, i));
   }
   return res;
 }
 
 SmallVector<Range, 4> LinalgOp::createLoopRanges(OpBuilder &b, Location loc) {
   AffineMap map = getLoopsToShapesMap();
   unsigned numDims = map.getNumDims(), numRes = map.getNumResults();
   auto viewSizes = createFlatListOfOperandDims(b, loc);
   SmallVector<Range, 4> res(numDims);
   Value zeroVal = b.create<ConstantIndexOp>(loc, 0);
   Value oneVal = b.create<ConstantIndexOp>(loc, 1);
   for (unsigned idx = 0; idx < numRes; ++idx) {
     auto result = map.getResult(idx);
     if (auto d = result.dyn_cast<AffineDimExpr>()) {
       if (res[d.getPosition()].offset)
         continue;
       res[d.getPosition()] = Range{zeroVal, viewSizes[idx], oneVal};
     }
   }
struct HasAffineDimExprVisitor
   bool visitConstantExpr(AffineConstantExpr constExpr) { return false; }
 
   bool visitSymbolExpr(AffineSymbolExpr symbolExpr) { return false; }
 
 private:
   llvm::SmallSet<unsigned, 4> positions;
 };
 
-Optional<Value> LinalgOp::inferResultDimFromInputShapes(OpBuilder &b,
-                                                        Location loc,
-                                                        unsigned resultIdx,
-                                                        unsigned dim) {
+LogicalResult LinalgOp::reifyReturnTypeShapesPerResultDim(
+    OpBuilder &b, SmallVectorImpl<SmallVector<Value>> &reifiedReturnShapes) {
   // An example that helps understand the logic below.
   // Consider the following expression O(i+j, j) += A(i,k) * B(k, j)
   // We want to express the shape of dim 0 of O in terms of shape of the inputs.
   // This is achieved as follows.
   //   loopsToShapesMap = (d0, d1, d2) -> (d0, d2, d2, d1, d0 + d1, d1)
-  //   subMapOfResultDim = (d0, d1, d2) -> (d0 + d1)
+  //   subMapOfResultShapes = (d0, d1, d2) -> (d0 + d1, d1)
   //   shapesToLoopsMap = (d0, d2, d2, d3, d4, d5) -> (d0, d3, d2)
-  //   resultFromFromInputDim = subMapOfResultDim.compose(shapesToLoopMap)
-  //     = (d0, d1, d2, d3, d4, d5) -> (d0 + d1)
+  //   resultShapesFromInputShapes = subMapOfResultDim.compose(shapesToLoopMap)
+  //     = (d0, d1, d2, d3, d4, d5) -> (d0 + d1, d1)
   AffineMap loopsToShapesMap = getLoopsToShapesMap();
 
   // Find the position in the above map that represents the shape of the
   // result:dim being inferred.
-  Optional<unsigned> resultDimSubMapPos =
-      getResultValueDimPositionInLoopsToShapeMap(resultIdx, dim);
-  if (!resultDimSubMapPos)
-    return {};
+  auto resultShapesSubMapPos = getResultsPositionInLoopsToShapeMap();
 
   /// From loopsToShapesMap extract the submap that represents the shape of the
-  /// (resultIdx, dim) needed
-  AffineMap loopToResultDimShapeMap =
-      loopsToShapesMap.getSubMap(*resultDimSubMapPos);
-  AffineMap operandShapesToResultDimMap =
-      loopToResultDimShapeMap.compose(getShapesToLoopsMap());
+  /// (resultIdx, dim) needed.
+  SmallVector<unsigned, 4> resultPosRange =
+      llvm::to_vector<4>(llvm::seq<unsigned>(resultShapesSubMapPos.first,
+                                             resultShapesSubMapPos.second));
+  AffineMap loopToResultsShapeMap = loopsToShapesMap.getSubMap(resultPosRange);
+  AffineMap resultShapesFromInputShapesMap =
+      loopToResultsShapeMap.compose(getShapesToLoopsMap());
 
   // Check that the result dim map does not contain the positions corresponding
   // to the outputs.
   llvm::SmallSet<unsigned, 4> outputDims;
-  unsigned outputDimPosStart =
-      getResultValueDimPositionInLoopsToShapeMap(0, 0).getValue();
-  unsigned outputDimPosEnd =
-      getResultValueDimPositionInLoopsToShapeMap(getNumOutputs() - 1,
-                                                 getOutputOpOperands()
-                                                         .back()
-                                                         .get()
-                                                         .getType()
-                                                         .cast<ShapedType>()
-                                                         .getRank() -
-                                                     1)
-          .getValue();
-  llvm::for_each(llvm::seq<unsigned>(outputDimPosStart, outputDimPosEnd),
+  llvm::for_each(resultPosRange,
                  [&outputDims](unsigned dim) { outputDims.insert(dim); });
   HasAffineDimExprVisitor checkDimExpr(outputDims);
-  if (checkDimExpr.visit(operandShapesToResultDimMap.getResult(0)))
-    return llvm::None;
-  return applyMapToValues(b, loc, operandShapesToResultDimMap,
-                          createFlatListOfOperandDims(b, loc))[0];
+  Location loc = getOperation()->getLoc();
+  auto allResultDimValues =
+      applyMapToValues(b, loc, resultShapesFromInputShapesMap,
+                       createFlatListOfOperandDims(b, loc));
+  unsigned pos = 0;
+  ArrayRef<AffineExpr> shapeExprs = resultShapesFromInputShapesMap.getResults();
+  for (auto resultIdx : llvm::seq<unsigned>(0, getNumOutputs())) {
+    ShapedType resultType = getOutputShapedType(resultIdx);
+    SmallVector<Value> shapes;
+    for (unsigned dim : llvm::seq<unsigned>(0, resultType.getRank())) {
+      if (checkDimExpr.visit(shapeExprs[pos]))
+        shapes.push_back(
+            b.createOrFold<memref::DimOp>(loc, getOutput(resultIdx), dim));
+      else
+        shapes.push_back(allResultDimValues[pos]);
+      pos++;
+    }
+    reifiedReturnShapes.emplace_back(std::move(shapes));
+  }
+  return success();
 }
 
 LogicalResult mlir::linalg::detail::verifyStructuredOpInterface(Operation *op) {
   LinalgOp linalgOp = cast<LinalgOp>(op);
   // Expect at least one shaped operand.
   // This means an op that constructs a tensor out of indices cannot be a
   // LinalgOp at the moment. For now this will have to be a special op until we
   // have output shape operands that are not tensors.

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

 #include "mlir/Dialect/Linalg/EDSC/Intrinsics.h"
 #include "mlir/Dialect/Linalg/IR/LinalgTypes.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/IR/AffineExprVisitor.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/PatternMatch.h"
+#include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Parser.h"
 
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/StringSet.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/MathExtras.h"
parseNamedStructuredOp(OpAsmParser &parser, OperationState &result,
                        ArrayRef<OpAsmParser::OperandType> captures = {});
 
 static void printNamedStructuredOpResults(OpAsmPrinter &p,
                                           TypeRange resultTypes);
 
 template <typename NamedStructuredOpType>
 static void printNamedStructuredOp(OpAsmPrinter &p, NamedStructuredOpType op);
 
+/// Helper function to convert a Value into an OpFoldResult, if the Value is
+/// known to be a constant index value.
+static SmallVector<OpFoldResult> getAsOpFoldResult(ArrayRef<Value> values) {
+  return llvm::to_vector<4>(
+      llvm::map_range(values, [](Value v) -> OpFoldResult {
+        APInt intValue;
+        if (v.getType().isa<IndexType>() &&
+            matchPattern(v, m_ConstantInt(&intValue))) {
+          return IntegerAttr::get(v.getType(), intValue.getSExtValue());
+        }
+        return v;
+      }));
+}
+
+/// Helper function to convert a vector of `OpFoldResult`s into a vector of
+/// `Value`s.
+static SmallVector<Value> getAsValues(OpBuilder &b, Location loc,
+                                      ArrayRef<OpFoldResult> valueOrAttrVec) {
+  return llvm::to_vector<4>(
+      llvm::map_range(valueOrAttrVec, [&](OpFoldResult value) -> Value {
+        if (auto attr = value.dyn_cast<Attribute>())
+          return b.create<ConstantIndexOp>(loc,
+                                           attr.cast<IntegerAttr>().getInt());
+        return value.get<Value>();
+      }));
+}
+
 /// Helper function to dispatch an OpFoldResult into either the `dynamicVec` if
 /// it is a Value or into `staticVec` if it is an IntegerAttr.
 /// In the case of a Value, a copy of the `sentinel` value is also pushed to
 /// `staticVec`. This is useful to extract mixed static and dynamic entries that
 /// come from an AttrSizedOperandSegments trait.
 static void dispatchIndexOpFoldResult(OpFoldResult ofr,
                                       SmallVectorImpl<Value> &dynamicVec,
                                       SmallVectorImpl<int64_t> &staticVec,

 //===----------------------------------------------------------------------===//
 void InitTensorOp::build(OpBuilder &b, OperationState &result,
                          ArrayRef<OpFoldResult> sizes, Type elementType,
                          ArrayRef<NamedAttribute> attrs) {
   unsigned rank = sizes.size();
   SmallVector<Value, 4> dynamicSizes;
   SmallVector<int64_t, 4> staticSizes;
   for (unsigned i = 0; i < rank; ++i) {
-    // staticLow and staticHigh have full information of the padding config.
-    // This will grow staticLow and staticHigh with 1 value. If the config is
-    // dynamic (ie not a constant), dynamicLow and dynamicHigh will grow with 1
-    // value as well.
     dispatchIndexOpFoldResult(sizes[i], dynamicSizes, staticSizes,
                               ShapedType::kDynamicSize);
   }
   auto resultType = RankedTensorType ::get(staticSizes, elementType);
   build(b, result, resultType, dynamicSizes, b.getI64ArrayAttr(staticSizes));
   result.addAttributes(attrs);
 }
 
if (newType == op.getType())
       return failure();
     auto newOp =
         rewriter.create<InitTensorOp>(op.getLoc(), newType, dynamicSizes,
                                       rewriter.getI64ArrayAttr(staticSizes));
     rewriter.replaceOpWithNewOp<tensor::CastOp>(op, op.getType(), newOp);
     return success();
   }
 };
-
-/// Canonicalize a `linalg.init_tensor` -> `dim` pattern by replacing the `dim`
-/// with
-/// - A constant value if the size is static along the dimension.
-/// - The dynamic value that defines the size of the result of
-///   `linalg.init_tensor` op.
-struct ReplaceDimOfInitTensorOp : public OpRewritePattern<memref::DimOp> {
-  using OpRewritePattern<memref::DimOp>::OpRewritePattern;
-
-  LogicalResult matchAndRewrite(memref::DimOp dimOp,
-                                PatternRewriter &rewriter) const override {
-    auto initTensorOp = dimOp.memrefOrTensor().getDefiningOp<InitTensorOp>();
-    if (!initTensorOp)
-      return failure();
-    auto dimIndex = dimOp.index().getDefiningOp<ConstantIndexOp>();
-    if (!dimIndex)
-      return failure();
-    int64_t index = dimIndex.getValue();
-    if (!initTensorOp.isDynamicSize(index)) {
-      rewriter.replaceOpWithNewOp<ConstantIndexOp>(
-          dimOp, initTensorOp.getStaticSize(index));
-    } else {
-      rewriter.replaceOp(dimOp, initTensorOp.getDynamicSize(index));
-    }
-    return success();
-  }
-};
 } // namespace
 
 namespace {
 /// Since `init_tensor` operation creates a tensor needed only for its shape, a
 /// subtensor of this is also needed only for its shape. The result can be
 /// replaced by a new init_tensor operation of the same size as the subtensor
 /// op.
 struct FoldInitTensorWithSubTensorOp : public OpRewritePattern<SubTensorOp> {
struct FoldInitTensorWithTensorReshapeOp
     : public OpRewritePattern<TensorReshapeOp> {
   using OpRewritePattern<TensorReshapeOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(TensorReshapeOp reshapeOp,
                                 PatternRewriter &rewriter) const override {
     if (!reshapeOp.src().getDefiningOp<InitTensorOp>())
       return failure();
     Location loc = reshapeOp.getLoc();
-    SmallVector<Value, 4> resultShapeValues =
-        reshapeOp.getOutputShape(rewriter, loc);
+    SmallVector<SmallVector<Value>, 4> resultShapes;
+    if (failed(reshapeOp.reifyReturnTypeShapesPerResultDim(rewriter,
+                                                           resultShapes)) ||
+        !llvm::hasSingleElement(resultShapes))
+      return failure();
     Value initTensor = rewriter.create<InitTensorOp>(
-        loc, resultShapeValues, reshapeOp.getResultType().getElementType());
+        loc, getAsOpFoldResult(resultShapes[0]),
+        reshapeOp.getResultType().getElementType());
+    if (initTensor.getType() != reshapeOp.getResultType()) {
       rewriter.replaceOpWithNewOp<tensor::CastOp>(
           reshapeOp, reshapeOp.getResultType(), initTensor);
+    } else {
+      rewriter.replaceOp(reshapeOp, initTensor);
+    }
     return success();
   }
 };
 } // namespace
 
 void InitTensorOp::getCanonicalizationPatterns(
     OwningRewritePatternList &results, MLIRContext *context) {
-  results
-      .insert<FoldInitTensorWithSubTensorOp, FoldInitTensorWithTensorReshapeOp,
-              ReplaceDimOfInitTensorOp, ReplaceStaticShapeDims>(context);
+  results.insert<FoldInitTensorWithSubTensorOp,
+                 FoldInitTensorWithTensorReshapeOp, ReplaceStaticShapeDims>(
+      context);
+}
+
+LogicalResult InitTensorOp::reifyReturnTypeShapesPerResultDim(
+    OpBuilder &builder,
+    SmallVectorImpl<SmallVector<Value>> &reifiedReturnShapes) {
+  auto shapes = llvm::to_vector<4>(llvm::map_range(
+      llvm::seq<int64_t>(0, getType().getRank()), [&](int64_t dim) -> Value {
+        if (isDynamicSize(dim))
+          return getDynamicSize(dim);
+        return builder.create<ConstantIndexOp>(getLoc(), getStaticSize(dim));
+      }));
+  reifiedReturnShapes.emplace_back(std::move(shapes));
+  return success();
 }
 
 //===----------------------------------------------------------------------===//
 // PadTensorOp
 //===----------------------------------------------------------------------===//
 
 /// Extract int64_t values from the assumed ArrayAttr of IntegerAttr.
 static SmallVector<int64_t, 4> extractFromI64ArrayAttr(Attribute attr) {
for (int i = 0; i < rank; ++i) {
     auto highValue = builder.createOrFold<SubIOp>(loc, resultDimSize, dimOp);
     high.push_back(highValue);
     low.push_back(builder.createOrFold<ConstantIndexOp>(loc, 0));
   }
   return PadTensorOp::createPadScalarOp(type, source, pad, low, high, loc,
                                         builder);
 }
 
+LogicalResult PadTensorOp::reifyReturnTypeShapesPerResultDim(
+    OpBuilder &b, SmallVectorImpl<SmallVector<Value>> &reifiedReturnShapes) {
+  Location loc = getLoc();
+  auto lowPad = getMixedLowPad();
+  auto highPad = getMixedHighPad();
+  SmallVector<Value> shapes;
+  for (auto dim : llvm::seq<int64_t>(0, getSourceType().getRank())) {
+    // Shape along each dimension is source dim + low pad + high pad.
+    SmallVector<Value> mapOperands;
+    mapOperands.push_back(b.createOrFold<memref::DimOp>(loc, source(), dim));
+    AffineExpr expr = b.getAffineDimExpr(0);
+    unsigned numSymbols = 0;
+    auto addOpFoldResult = [&](OpFoldResult valueOrAttr) {
+      if (Value v = valueOrAttr.dyn_cast<Value>()) {
+        expr = expr + b.getAffineSymbolExpr(numSymbols++);
+        mapOperands.push_back(v);
+        return;
+      }
+      int64_t staticValue =
+          valueOrAttr.get<Attribute>().cast<IntegerAttr>().getInt();
+      expr = expr + staticValue;
+    };
+    addOpFoldResult(lowPad[dim]);
+    addOpFoldResult(highPad[dim]);
+    shapes.push_back(applyMapToValues(
+        b, loc, AffineMap::get(1, numSymbols, expr), mapOperands)[0]);
+  }
+  reifiedReturnShapes.emplace_back(std::move(shapes));
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // ReshapeOp
 //===----------------------------------------------------------------------===//
 
 /// Collapse reassociation maps that are used in pair of reshape ops where one
 /// is a producer and other is the consumer. Only valid to use this method when
 /// both the producer and consumer are collapsing dimensions or both are
 /// expanding dimensions.
convertReassociationIndicesToMaps(
   }
   return reassociationMaps;
 }
 
 /// For reshape op compute the shape at dimension `dimIndex` of the output in
 /// terms of shape of the `src`, when the reshape op is a collapsing
 /// operation. It is the product of the shape of the collapsed dimensions of the
 /// `src`.
-static Value
+static OpFoldResult
 getCollapsedOutputDimFromInputShape(OpBuilder &builder, Location loc,
                                     int64_t dimIndex, Value src,
                                     ArrayRef<AffineMap> reassociationMap) {
   AffineMap map = reassociationMap[dimIndex];
   unsigned startPos =
       map.getResults().front().cast<AffineDimExpr>().getPosition();
   unsigned endPos = map.getResults().back().cast<AffineDimExpr>().getPosition();
   AffineExpr expr;
   SmallVector<Value, 2> dynamicDims;
   for (auto dim : llvm::seq(startPos, endPos + 1)) {
-    dynamicDims.push_back(builder.create<memref::DimOp>(loc, src, dim));
+    dynamicDims.push_back(builder.createOrFold<memref::DimOp>(loc, src, dim));
     AffineExpr currExpr = builder.getAffineSymbolExpr(dim - startPos);
     expr = (expr ? expr * currExpr : currExpr);
   }
   return applyMapToValues(builder, loc,
                           AffineMap::get(0, endPos - startPos + 1, expr),
                           dynamicDims)[0];
 }
 
 /// Given the `src` of a collapsing reshape op and its reassociation maps,
 /// compute the shape of the result of the reshape.
-static SmallVector<Value, 4> getCollapsedOutputShapeFromInputShape(
+static SmallVector<OpFoldResult, 4> getCollapsedOutputShapeFromInputShape(
     OpBuilder &builder, Location loc, Value src,
     ArrayRef<int64_t> dstStaticShape, ArrayRef<AffineMap> reassociation) {
   return llvm::to_vector<4>(llvm::map_range(
       llvm::seq<int64_t>(0, dstStaticShape.size()), [&](int64_t dim) {
         return getCollapsedOutputDimFromInputShape(builder, loc, dim, src,
                                                    reassociation);
       }));
 }
for (auto dim : llvm::seq(startPos, endPos + 1)) {
       expandedDimToCollapsedDim[dim] = map.index();
     }
   }
   return expandedDimToCollapsedDim;
 }
 
 /// For an expanding reshape op, compute the value for a dimension of the output
 /// from the shape of the input.
-static Value getExpandedOutputDimFromInputShape(
+static OpFoldResult getExpandedOutputDimFromInputShape(
     OpBuilder &builder, Location loc, int64_t dimIndex, Value src,
     ArrayRef<int64_t> dstStaticShape, ArrayRef<AffineMap> reassociation,
     llvm::DenseMap<int64_t, int64_t> &expandedDimToCollapsedDim) {
   if (!ShapedType::isDynamic(dstStaticShape[dimIndex])) {
-    return builder.create<ConstantIndexOp>(loc, dstStaticShape[dimIndex]);
+    return builder.getI64IntegerAttr(dstStaticShape[dimIndex]);
   }
   unsigned sourceDimPos = expandedDimToCollapsedDim[dimIndex];
   unsigned startPos = reassociation[sourceDimPos]
                           .getResults()
                           .front()
                           .cast<AffineDimExpr>()
                           .getPosition();
   unsigned endPos = reassociation[sourceDimPos]
return applyMapToValues(
       builder, loc,
       AffineMap::get(
           0, 1, builder.getAffineSymbolExpr(0).floorDiv(linearizedStaticDim)),
       sourceDim)[0];
 }
 
 /// Given the `src` of an expanding reshape op, the reassociation maps and the
 /// result type, compute the shape of the result of the reshape.
-static SmallVector<Value, 4> getExpandedOutputShapeFromInputShape(
+static SmallVector<OpFoldResult, 4> getExpandedOutputShapeFromInputShape(
     OpBuilder &builder, Location loc, Value src,
     ArrayRef<int64_t> dstStaticShape, ArrayRef<AffineMap> reassociation) {
   llvm::DenseMap<int64_t, int64_t> expandedDimToCollapsedDim =
       getExpandedDimToCollapsedDimMap(reassociation);
   return llvm::to_vector<4>(llvm::map_range(
       llvm::seq<int64_t>(0, dstStaticShape.size()), [&](int64_t dim) {
         return getExpandedOutputDimFromInputShape(builder, loc, dim, src,
                                                   dstStaticShape, reassociation,
                                                   expandedDimToCollapsedDim);
       }));
 }
 
-SmallVector<Value, 4> mlir::linalg::getReshapeOutputShapeFromInputShape(
-    OpBuilder &builder, Location loc, Value src,
-    ArrayRef<int64_t> dstStaticShape, ArrayRef<AffineMap> reassocation) {
+static SmallVector<OpFoldResult, 4>
+getReshapeOutputShapeFromInputShape(OpBuilder &builder, Location loc, Value src,
+                                    ArrayRef<int64_t> dstStaticShape,
+                                    ArrayRef<AffineMap> reassocation) {
   return dstStaticShape.size() >
                  static_cast<size_t>(src.getType().cast<ShapedType>().getRank())
              ? getExpandedOutputShapeFromInputShape(
                    builder, loc, src, dstStaticShape, reassocation)
              : getCollapsedOutputShapeFromInputShape(
                    builder, loc, src, dstStaticShape, reassocation);
 }
 
-/// For a reshape op, compute the value of a given dimension of the output
-/// (`dimIndex`) from the shape of the inputs and type of the result.
-static Value getReshapeOutputDimFromInputShape(
-    OpBuilder &builder, Location loc, int64_t dimIndex, Value src,
-    ArrayRef<int64_t> dstStaticShape, ArrayRef<AffineMap> reassociation) {
-  if (dstStaticShape.size() >
-      static_cast<size_t>(src.getType().cast<ShapedType>().getRank())) {
-    llvm::DenseMap<int64_t, int64_t> expandedDimToCollapsedDim =
-        getExpandedDimToCollapsedDimMap(reassociation);
-    return getExpandedOutputDimFromInputShape(builder, loc, dimIndex, src,
-                                              dstStaticShape, reassociation,
-                                              expandedDimToCollapsedDim);
-  }
-  return getCollapsedOutputDimFromInputShape(builder, loc, dimIndex, src,
-                                             reassociation);
-}
-
 void mlir::linalg::ReshapeOp::build(OpBuilder &b, OperationState &result,
                                     Value src,
                                     ArrayRef<ReassociationExprs> reassociation,
                                     ArrayRef<NamedAttribute> attrs) {
   auto maps = getSymbolLessAffineMaps(reassociation);
   auto memRefType = src.getType().cast<MemRefType>();
   auto resultType = computeReshapeCollapsedType(memRefType, maps);
   build(b, result, resultType, src, attrs);
LogicalResult matchAndRewrite(TensorReshapeOp reshapeOp,
     if (!attr || !attr.isSplat())
       return failure();
     DenseElementsAttr newAttr = DenseElementsAttr::getFromRawBuffer(
         reshapeOp.getResultType(), attr.getRawData(), true);
     rewriter.replaceOpWithNewOp<ConstantOp>(reshapeOp, newAttr);
     return success();
   }
 };
-
-/// Canonicalize dim ops that use the output shape with dim of the input.
-struct ReplaceDimOfReshapeOpResult : OpRewritePattern<memref::DimOp> {
-  using OpRewritePattern<memref::DimOp>::OpRewritePattern;
-  LogicalResult matchAndRewrite(memref::DimOp dimOp,
-                                PatternRewriter &rewriter) const override {
-    Value dimValue = dimOp.memrefOrTensor();
-    Optional<int64_t> dimIndex = dimOp.getConstantIndex();
-    if (!dimIndex)
-      return failure();
-
-    auto reshapeOp = dimValue.getDefiningOp<TensorReshapeOp>();
-    if (!reshapeOp)
-      return failure();
-
-    rewriter.replaceOp(dimOp,
-                       getReshapeOutputDimFromInputShape(
-                           rewriter, dimOp.getLoc(), *dimIndex, reshapeOp.src(),
-                           reshapeOp.getResultType().getShape(),
-                           reshapeOp.getReassociationMaps()));
-    return success();
-  }
-};
 } // namespace
 
 void TensorReshapeOp::getCanonicalizationPatterns(
     OwningRewritePatternList &results, MLIRContext *context) {
-  results.insert<CollapseReshapeOps<TensorReshapeOp>, FoldReshapeWithConstant,
-                 ReplaceDimOfReshapeOpResult>(context);
+  results.insert<CollapseReshapeOps<TensorReshapeOp>, FoldReshapeWithConstant>(
+      context);
+}
+
+LogicalResult TensorReshapeOp::reifyReturnTypeShapesPerResultDim(
+    OpBuilder &b, SmallVectorImpl<SmallVector<Value>> &reifiedReturnShapes) {
+  auto resultShape =
+      getAsValues(b, getLoc(),
+                  getReshapeOutputShapeFromInputShape(
+                      b, getLoc(), src(), getResultType().getShape(),
+                      getReassociationMaps()));
+  reifiedReturnShapes.emplace_back(std::move(resultShape));
+  return success();
 }
 
 //===----------------------------------------------------------------------===//
 // YieldOp
 //===----------------------------------------------------------------------===//
 
 static void print(OpAsmPrinter &p, linalg::YieldOp op) {
   p << op.getOperationName();
for (auto result : llvm::zip(op->getResults(), newOp->getResults())) {
         replacements.push_back(newResult);
       }
     }
     rewriter.replaceOp(op, replacements);
 
     return success();
   }
 };
-
-/// Replaces memref.dim operations that use the result of a LinalgOp (on
-/// tensors) with memref.dim operations that use one of the arguments. For
-/// example,
-///
-/// %0 = linalg.matmul ins(%arg0, %arg1, ...)
-/// %1 = memref.dim %0, %c0
-///
-/// with
-///
-/// %1 = memref.dim %arg0, %c0
-///
-/// where possible. With this the result of the `linalg.matmul` is not used in
-/// dim operations. If the value produced is replaced with another value (say by
-/// tiling `linalg.matmul`) will make the `linalg.matmul` truly dead instead of
-/// used in a dim op that would prevent the DCE of this op.
-struct ReplaceDimOfLinalgOpResult : public OpRewritePattern<memref::DimOp> {
-  using OpRewritePattern<memref::DimOp>::OpRewritePattern;
-
-  LogicalResult matchAndRewrite(memref::DimOp dimOp,
-                                PatternRewriter &rewriter) const override {
-    Value dimValue = dimOp.memrefOrTensor();
-    Optional<int64_t> dimIndex = dimOp.getConstantIndex();
-    if (!dimIndex)
-      return failure();
-    auto linalgOp = dimValue.getDefiningOp<LinalgOp>();
-    if (!linalgOp)
-      return failure();
-
-    unsigned resultIndex = dimValue.cast<OpResult>().getResultNumber();
-    Optional<Value> operandDimValue = linalgOp.inferResultDimFromInputShapes(
-        rewriter, dimOp.getLoc(), resultIndex,
-        static_cast<unsigned>(*dimIndex));
-    if (!operandDimValue) {
-      // Its always possible to replace using the corresponding `outs`
-      // parameter.
-      operandDimValue = rewriter.create<memref::DimOp>(
-          dimOp.getLoc(), linalgOp.getOutput(resultIndex), *dimIndex);
-    }
-    rewriter.replaceOp(dimOp, *operandDimValue);
-    return success();
-  }
-};
-
 } // namespace
 
 namespace {
 // Deduplicate redundant args of a linalg op.
 // An arg is redundant if it has the same Value and indexing map as another.
 struct DeduplicateInputs : public RewritePattern {
   DeduplicateInputs(PatternBenefit benefit = 1)
       : RewritePattern(benefit, MatchAnyOpTypeTag()) {}

 };
 } // namespace
 
 #define CANONICALIZERS_AND_FOLDERS(XXX)                                        \
   void XXX::getCanonicalizationPatterns(OwningRewritePatternList &results,     \
                                         MLIRContext *context) {                \
     results.insert<DeduplicateInputs, EraseDeadLinalgOp, FoldTensorCastOp,     \
                    RemoveIdentityLinalgOps>();                                 \
-    results.insert<ReplaceDimOfLinalgOpResult>(context);                       \
   }                                                                            \
                                                                                \
   LogicalResult XXX::fold(ArrayRef<Attribute>,                                 \
                           SmallVectorImpl<OpFoldResult> &) {                   \
     return foldMemRefCast(*this);                                              \
   }
 
 CANONICALIZERS_AND_FOLDERS(ConvOp)

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

 #include "mlir/Dialect/StandardOps/Utils/Utils.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeUtilities.h"
+#include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "llvm/ADT/STLExtras.h"
 
 using namespace mlir;
 using namespace mlir::memref;
 
 /// Materialize a single constant operation from a given attribute value with
 /// the desired resultant type.
 Operation *MemRefDialect::materializeConstant(OpBuilder &builder,
LogicalResult matchAndRewrite(DimOp dimOp,
     auto castOp = dimOp.memrefOrTensor().getDefiningOp<CastOpTy>();
     if (!castOp)
       return failure();
     Value newSource = castOp.getOperand();
     rewriter.replaceOpWithNewOp<DimOp>(dimOp, newSource, dimOp.index());
     return success();
   }
 };
+
+/// Helper method to get the `Value` that is the shape of the `resultIdx`-th
+/// result at dimension `dimIndex` from the `ShapedTypeOpInterface`.
+/// TODO(ravishankarm): This is better put as a interface utility method
+/// somewhere, but that would imply the interface will depend on the `tensor`
+/// dialect. Ideally maybe a utility method in the `tensor` dialect.
+static Value getResultDimFromShapeInterface(OpBuilder &builder, OpResult result,
+                                            int64_t dimIndex) {
+  unsigned resultNumber = result.getResultNumber();
+  auto shapedTypeOp = dyn_cast<InferShapedTypeOpInterface>(result.getOwner());
+  Location loc = result.getOwner()->getLoc();
+  if (!shapedTypeOp)
+    return nullptr;
+
+  // The interface exposes two methods, one that returns the shape of all the
+  // results as `Value` and other that returns the shape as a list of
+  // `SmallVector<Value>`. The former takes precedence over the latter. So first
+  // check if the op implements the first interface method or the second, and
+  // get the value to use appropriately.
+  SmallVector<Value> reifiedResultShapes;
+  if (succeeded(
+          shapedTypeOp.reifyReturnTypeShapes(builder, reifiedResultShapes))) {
+    if (reifiedResultShapes.size() <= resultNumber)
+      return nullptr;
+    Value resultShape = reifiedResultShapes[resultNumber];
+    auto resultShapeType = resultShape.getType().dyn_cast<RankedTensorType>();
+    if (!resultShapeType || !resultShapeType.getElementType().isa<IndexType>())
+      return nullptr;
+    return builder.create<tensor::ExtractOp>(
+        loc, resultShape, builder.createOrFold<ConstantIndexOp>(loc, dimIndex));
+  }
+
+  SmallVector<SmallVector<Value>> reifiedResultShapesPerDim;
+  if (failed(shapedTypeOp.reifyReturnTypeShapesPerResultDim(
+          builder, reifiedResultShapesPerDim)))
+    return nullptr;
+  if (reifiedResultShapesPerDim.size() <= resultNumber ||
+      reifiedResultShapesPerDim[resultNumber].size() !=
+          static_cast<size_t>(result.getType().cast<ShapedType>().getRank()))
+    return nullptr;
+  OpFoldResult valueOrAttr = reifiedResultShapesPerDim[resultNumber][dimIndex];
+  if (auto attr = valueOrAttr.dyn_cast<Attribute>())
+    return builder.createOrFold<ConstantIndexOp>(
+        loc, attr.cast<IntegerAttr>().getInt());
+  return valueOrAttr.get<Value>();
+}
+
+/// Fold dim of an operation that implements the InferShapedTypeOpInterface
+struct DimOfShapedTypeOpInterface : public OpRewritePattern<DimOp> {
+  using OpRewritePattern<DimOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(DimOp dimOp,
+                                PatternRewriter &rewriter) const override {
+    OpResult dimValue = dimOp.memrefOrTensor().dyn_cast<OpResult>();
+    if (!dimValue)
+      return failure();
+    auto shapedTypeOp =
+        dyn_cast<InferShapedTypeOpInterface>(dimValue.getOwner());
+    if (!shapedTypeOp)
+      return failure();
+
+    Optional<int64_t> dimIndex = dimOp.getConstantIndex();
+    if (!dimIndex)
+      return failure();
+    Value replacement =
+        getResultDimFromShapeInterface(rewriter, dimValue, *dimIndex);
+    if (!replacement)
+      return failure();
+    rewriter.replaceOp(dimOp, replacement);
+    return success();
+  }
+};
 } // end anonymous namespace.
 
 void DimOp::getCanonicalizationPatterns(OwningRewritePatternList &results,
                                         MLIRContext *context) {
   results.insert<DimOfMemRefReshape, DimOfCastOp<BufferCastOp>,
-                 DimOfCastOp<tensor::CastOp>>(context);
+                 DimOfCastOp<tensor::CastOp>, DimOfShapedTypeOpInterface>(
+      context);
 }
 
 // ---------------------------------------------------------------------------
 // DmaStartOp
 // ---------------------------------------------------------------------------
 
 void DmaStartOp::build(OpBuilder &builder, OperationState &result,
                        Value srcMemRef, ValueRange srcIndices, Value destMemRef,

mlir/lib/Interfaces/InferTypeOpInterface.cpp

 //===- InferTypeOpInterface.cpp - Infer Type Interfaces ---------*- 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file contains the definitions of the infer op interfaces defined in
 // `InferTypeOpInterface.td`.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Interfaces/InferTypeOpInterface.h"
-
 #include "mlir/IR/BuiltinTypes.h"
 
 using namespace mlir;
 
 namespace mlir {
 #include "mlir/Interfaces/InferTypeOpInterface.cpp.inc"
 } // namespace mlir
 
 LogicalResult mlir::detail::inferReturnTensorTypes(
     function_ref<LogicalResult(

mlir/test/Dialect/Linalg/canonicalize.mlir

 // CHECK-SAME:   %[[ARG0:[a-zA-Z0-9_]+]]: index
 // CHECK-SAME:   %[[ARG1:[a-zA-Z0-9_]+]]: index
 //      CHECK:   return %[[ARG0]], %[[ARG1]]
 
 // -----
 
 func @remove_dim_result_uses
   (%arg0 : tensor<?x?xf32>, %arg1 : tensor<?x?xf32>,
-   %arg2 : tensor<?x?xf32>) -> (index) {
+   %arg2 : tensor<?x?xf32>) -> (index, index) {
   %c0 = constant 0 : index
+  %c1 = constant 1 : index
   %0 = linalg.generic
     {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d2)>,
                       affine_map<(d0, d1, d2) -> (d2, d1)>,
-                      affine_map<(d0, d1, d2) -> (d0 + d1, d1)>],
+                      affine_map<(d0, d1, d2) -> (d0 + d1, d1 - d0)>],
      iterator_types = ["parallel", "parallel", "reduction"]}
     ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>)
     outs(%arg2 : tensor<?x?xf32>) {
     ^bb0(%arg3 : f32, %arg4 : f32, %arg5 : f32):
       %1 = mulf %arg3, %arg4 : f32
       %2 = addf %1, %arg5 : f32
       linalg.yield %2 : f32
     } -> tensor<?x?xf32>
   %3 = memref.dim %0, %c0 : tensor<?x?xf32>
-  return %3 : index
+  %4 = memref.dim %0, %c1 : tensor<?x?xf32>
+  return %3, %4 : index, index
 }
-//       CHECK: #[[MAP:.+]] = affine_map<()[s0, s1] -> (s0 + s1)>
+//       CHECK: #[[MAP0:.+]] = affine_map<()[s0, s1] -> (s0 + s1)>
+//       CHECK: #[[MAP1:.+]] = affine_map<()[s0, s1] -> (-s0 + s1)>
 //       CHECK: func @remove_dim_result_uses
 //  CHECK-SAME:   %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?x?xf32>
 //  CHECK-SAME:   %[[ARG1:[a-zA-Z0-9_]+]]: tensor<?x?xf32>
 //  CHECK-SAME:   %[[ARG2:[a-zA-Z0-9_]+]]: tensor<?x?xf32>
 //   CHECK-DAG:   %[[C0:.+]] = constant 0 : index
 //   CHECK-DAG:   %[[C1:.+]] = constant 1 : index
 //   CHECK-DAG:   %[[T0:.+]] = memref.dim %[[ARG0]], %[[C0]]
 //   CHECK-DAG:   %[[T1:.+]] = memref.dim %[[ARG1]], %[[C1]]
-//       CHECK:   %[[T2:.+]] = affine.apply #[[MAP]]()[%[[T0]], %[[T1]]]
-//       CHECK:   return %[[T2]]
+//       CHECK:   %[[T2:.+]] = affine.apply #[[MAP0]]()[%[[T0]], %[[T1]]]
+//   CHECK-DAG:   %[[T3:.+]] = memref.dim %[[ARG0]], %[[C0]]
+//   CHECK-DAG:   %[[T4:.+]] = memref.dim %[[ARG1]], %[[C1]]
+//       CHECK:   %[[T5:.+]] = affine.apply #[[MAP1]]()[%[[T3]], %[[T4]]]
+//       CHECK:   return %[[T2]], %[[T5]]
 
 // -----
 
 func @remove_dim_result_uses_outs
   (%arg0 : tensor<?xf32>, %arg1 : index) -> (index) {
   %c0 = constant 0 : index
   %c1 = constant 1 : index
   %d0 = memref.dim %arg0, %c0 : tensor<?xf32>

 func @self_copy(%arg0 : memref<2x3x?x4xf32>) {
 
 //   CHECK-NOT: linalg.copy
   linalg.copy(%arg0, %arg0): memref<2x3x?x4xf32>, memref<2x3x?x4xf32>
 
 //   CHECK: return
   return
 }
+
+// -----
+
+func @dim_of_pad_op(%arg0 : tensor<2x?x?xf32>, %arg1 : index, %arg2 : index,
+    %arg3: f32) -> (index, index, index)
+{
+   %c0 = constant 0 : index
+   %c1 = constant 1 : index
+   %c2 = constant 2 : index
+   %c3 = constant 3 : index
+   %c4 = constant 4 : index
+   %c5 = constant 5 : index
+   %0 = linalg.pad_tensor %arg0 low[%c3, %arg1, %c4] high[7, %c5, %arg2] {
+     ^bb0(%arg4: index, %arg5: index, %arg6: index):
+       linalg.yield %arg3 : f32
+   } : tensor<2x?x?xf32> to tensor<?x?x?xf32>
+   %1 = memref.dim %0, %c0 : tensor<?x?x?xf32>
+   %2 = memref.dim %0, %c1 : tensor<?x?x?xf32>
+   %3 = memref.dim %0, %c2 : tensor<?x?x?xf32>
+   return %1, %2, %3 : index, index, index
+}
+//  CHECK-DAG: #[[MAP0:.+]] = affine_map<()[s0, s1] -> (s0 + s1 + 5)>
+//  CHECK-DAG: #[[MAP1:.+]] = affine_map<()[s0, s1] -> (s0 + s1 + 4)>
+//      CHECK: func @dim_of_pad_op
+// CHECK-SAME:   %[[ARG0:[A-Za-z0-9_]+]]: tensor<2x?x?xf32>
+// CHECK-SAME:   %[[ARG1:[A-Za-z0-9_]+]]: index
+// CHECK-SAME:   %[[ARG2:[A-Za-z0-9_]+]]: index
+//  CHECK-DAG:   %[[C1:.+]] = constant 1 : index
+//  CHECK-DAG:   %[[C2:.+]] = constant 2 : index
+//  CHECK-DAG:   %[[C12:.+]] = constant 12 : index
+//      CHECK:   %[[IN_DIM1:.+]] = memref.dim %[[ARG0]], %[[C1]]
+//      CHECK:   %[[OUT_DIM1:.+]] = affine.apply #[[MAP0]]()[%[[ARG1]], %[[IN_DIM1]]]
+//      CHECK:   %[[IN_DIM2:.+]] = memref.dim %[[ARG0]], %[[C2]]
+//      CHECK:   %[[OUT_DIM2:.+]] = affine.apply #[[MAP1]]()[%[[ARG2]], %[[IN_DIM2]]]
+//      CHECK:   return %[[C12]], %[[OUT_DIM1]], %[[OUT_DIM2]]

mlir/test/Transforms/test-canonicalize.mlir

func @typemismatch() -> i32 {
 
   // The "passthrough_fold" folder will naively return its operand, but we don't
   // want to fold here because of the type mismatch.
 
   // CHECK: "test.passthrough_fold"
   %0 = "test.passthrough_fold"(%c42) : (f32) -> (i32)
   return %0 : i32
 }
+
+// CHECK-LABEL: func @result_shape_per_dim
+// CHECK-SAME: (%[[ARG_0:[a-z0-9]*]]: tensor<2x3x?xf32>, %[[ARG_1:[a-z0-9]*]]: tensor<?x5xf32>)
+func @result_shape_per_dim(%arg0 : tensor<2x3x?xf32>, %arg1 : tensor<?x5xf32>)
+    -> (index, index, index, index, index) {
+  // CHECK-DAG: %[[C0:.+]] = constant 0 : index
+  // CHECK-DAG: %[[C2:.+]] = constant 2 : index
+  // CHECK-DAG: %[[C3:.+]] = constant 3 : index
+  // CHECK-DAG: %[[C5:.+]] = constant 5 : index
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+  %c2 = constant 2 : index
+  %0:2 = "test.op_with_result_shape_per_dim_interface"(%arg0, %arg1)
+      : (tensor<2x3x?xf32>, tensor<?x5xf32>) -> (tensor<?x5xf32>, tensor<2x3x?xf32>)
+  %1 = memref.dim %0#0, %c0 : tensor<?x5xf32>
+  %2 = memref.dim %0#0, %c1 : tensor<?x5xf32>
+  %3 = memref.dim %0#1, %c0 : tensor<2x3x?xf32>
+  %4 = memref.dim %0#1, %c1 : tensor<2x3x?xf32>
+  %5 = memref.dim %0#1, %c2 : tensor<2x3x?xf32>
+  // CHECK-DAG: %[[D0:.+]] = memref.dim %[[ARG_1]], %[[C0]]
+  // CHECK-DAG: %[[D1:.+]] = memref.dim %[[ARG_0]], %[[C2]]
+  // CHECK: return %[[D0]], %[[C5]], %[[C2]], %[[C3]], %[[D1]]
+  return %1, %2, %3, %4, %5 : index, index, index, index, index
+}
 No newline at end of file

mlir/test/lib/Dialect/Test/TestDialect.cpp

 
 LogicalResult OpWithShapedTypeInferTypeInterfaceOp::reifyReturnTypeShapes(
     OpBuilder &builder, llvm::SmallVectorImpl<Value> &shapes) {
   shapes = SmallVector<Value, 1>{
       builder.createOrFold<memref::DimOp>(getLoc(), getOperand(0), 0)};
   return success();
 }
 
+LogicalResult
+OpWithResultShapePerDimInterfaceOp ::reifyReturnTypeShapesPerResultDim(
+    OpBuilder &builder,
+    llvm::SmallVectorImpl<llvm::SmallVector<Value>> &shapes) {
+  SmallVector<Value> operand1Shape, operand2Shape;
+  Location loc = getLoc();
+  for (auto i :
+       llvm::seq<int>(0, operand1().getType().cast<ShapedType>().getRank())) {
+    operand1Shape.push_back(builder.create<memref::DimOp>(loc, operand1(), i));
+  }
+  for (auto i :
+       llvm::seq<int>(0, operand2().getType().cast<ShapedType>().getRank())) {
+    operand2Shape.push_back(builder.create<memref::DimOp>(loc, operand2(), i));
+  }
+  shapes.emplace_back(std::move(operand2Shape));
+  shapes.emplace_back(std::move(operand1Shape));
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // Test SideEffect interfaces
 //===----------------------------------------------------------------------===//
 
 namespace {
 /// A test resource for side effects.
 struct TestResource : public SideEffects::Resource::Base<TestResource> {
   StringRef getName() final { return "<Test>"; }

mlir/test/lib/Dialect/Test/TestOps.td

def I32ElementsAttrOp : TEST_Op<"i32ElementsAttr"> {
   let arguments = (ins I32ElementsAttr:$attr);
 }
 
 def IndexElementsAttrOp : TEST_Op<"indexElementsAttr"> {
   let arguments = (ins IndexElementsAttr:$attr);
 }
 
 def OpWithInferTypeInterfaceOp : TEST_Op<"op_with_infer_type_if", [
-    DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
+    DeclareOpInterfaceMethods<InferTypeOpInterface,
+        ["inferReturnTypeComponents"]>]> {
   let arguments = (ins AnyTensor, AnyTensor);
   let results = (outs AnyTensor);
 }
 
 def OpWithShapedTypeInferTypeInterfaceOp : TEST_Op<"op_with_shaped_type_infer_type_if",
       InferTensorTypeWithReify.traits> {
   let arguments = (ins AnyTensor, AnyTensor);
   let results = (outs AnyTensor);
 }
 
+def OpWithResultShapePerDimInterfaceOp : TEST_Op<"op_with_result_shape_per_dim_interface",
+      [DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
+           ["reifyReturnTypeShapesPerResultDim"]>]> {
+  let arguments = (ins AnyRankedTensor:$operand1, AnyRankedTensor:$operand2);
+  let results = (outs AnyRankedTensor:$result1, AnyRankedTensor:$result2);
+}
+
 def IsNotScalar : Constraint<CPred<"$0.getType().getRank() != 0">>;
 
 def UpdateAttr : Pat<(I32ElementsAttrOp $attr),
                      (I32ElementsAttrOp ConstantAttr<I32ElementsAttr, "0">),
                      [(IsNotScalar $attr)]>;
 
 def TestBranchOp : TEST_Op<"br",
     [DeclareOpInterfaceMethods<BranchOpInterface>, Terminator]> {

mlir/test/lib/Dialect/Test/TestPatterns.cpp

 }
 
 static void reifyReturnShape(Operation *op) {
   OpBuilder b(op);
 
   // Use permutations of 2 args as operands.
   auto shapedOp = cast<OpWithShapedTypeInferTypeInterfaceOp>(op);
   SmallVector<Value, 2> shapes;
-  if (failed(shapedOp.reifyReturnTypeShapes(b, shapes)))
+  if (failed(shapedOp.reifyReturnTypeShapes(b, shapes)) ||
+      !llvm::hasSingleElement(shapes))
     return;
-  for (auto it : llvm::enumerate(shapes))
+  for (auto it : llvm::enumerate(shapes)) {
     op->emitRemark() << "value " << it.index() << ": "
                      << it.value().getDefiningOp();
   }
+}
 
 struct TestReturnTypeDriver
     : public PassWrapper<TestReturnTypeDriver, FunctionPass> {
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<memref::MemRefDialect>();
   }
 
   void runOnFunction() override {