llvm
Lower aten.view to tensor.reshape with support for most dynamic cases
@silvasean PTAL. I remember you commenting on a previous PR that tensor.reshape was not a good op to use here. Is that still the case?
For reference, here is the comment https://github.com/llvm/torch-mlir/pull/870#pullrequestreview-981594500
This was my concern as well, but some of the cases we are starting to see where we need support for aten.view were difficult with the previous approach so I was wondering if a hybrid approach that uses the ExpandShapeOp/CollapseShapeOp approach when possible and otherwise defaults to tensor.reshape would be good.
This was my concern as well, but some of the cases we are starting to see where we need support for aten.view were difficult with the previous approach so I was wondering if a hybrid approach that uses the ExpandShapeOp/CollapseShapeOp approach when possible and otherwise defaults to tensor.reshape would be good.
Can you create small IR unit tests for all the cases that need to be handled so we can look for a solution together? Generally the tensor.reshape op is really bad for transformations. I'd rather we just fail, as this will just manifest as a performance bug later and be much harder to track down vs having it just fail compilation.
I dumped the E2E cases into a file for IR unit tests, but without any filecheck stuff just so that we can see the IR. I can add the lowerings from this PR if that helps, but it seems like we are moving away from tensor.reshape.
The current implementation lowers to ExpandShapeOp and/or CollapseShapeOp depending on the form of the view op. Support for dynamic dims depends mostly on tensor.size ops allowing us to infer the relationship between the input and output shapes. Additionally inferred dimensions are only supported when we can compute them from static dims. One side note is that the current lowering assumes that dims that get the size from a tensor.size op are neither collapsed nor expanded, which is probably what a user will write but isn't in any way guaranteed.
I think there are three cases that we are looking to support now, although I could be missing something. For #1161 we need support for inferred dimensions with dynamic dims where all of the dims come from tensor.size ops. The changes to inferred dimension handling in this PR already handles this case but needs to be ported over to the Expand/Collapse implementation.
Second, we also need support for zero rank -> rank 1 but while this PR handles that with tensor.reshape I'm sure there are a number of ways to do that so we can just pick the best one.
The third case, is where either there is neither collapsing nor expanding, or the association of indices to collapse/expand is ambiguous ([2, 3] -> [3, 2] for static dims, and [?, ?] -> [?, ?, ?] for dynamic). My thinking is that the way to handle this is simply to flatten to an intermediate shape then expand back out ([2, 3] -> [6] -> [3, 2]) and thus we still end up with a single collapse and single expand op + some MulIOps and a cf.assert for the dynamic case. If we give up on improving past 1 expand + 1 collapse the question is how well we can do with minimizing the number of multiplications.
I think that most of the cases we actually need aren't too much of a problem with the 1 collapse + 1 expand approach, but I was thinking that the size of this lowering is just going to keep growing so I wanted to get an explanation as to why we weren't just using the tensor.reshape op that seemed to do what we want before jumping back into the current approach. A couple of the examples I'm thinking of are as follows (letters are dynamic but given by size ops): [a, b, c] -> [c, b, a] [a, b, c] -> [?, b, ?] (iirc this might actually be an incorrect lowering case currently based on the side note above, have to double check)
Can you create small IR unit tests for all the cases that need to be handled so we can look for a solution together? Generally the
tensor.reshapeop is really bad for transformations. I'd rather we just fail, as this will just manifest as a performance bug later and be much harder to track down vs having it just fail compilation.
Actually I'm still curious about something. If the tensor.reshape op is so bad for transformations, why not include a whole host of canonicalizations in mlir to avoid it as much as possible? Also what is it that makes is such a problematic op?
I think the collapse to 1d + expand approach is a good one that is quite general. We should only use it when there is no other way to handle the reshape.
@nicolasvasilache had a classification of reshapes, and I think ones like [2, 3] -> [3, 2] are the "bad case" where we can't do better than flattening and expanding. Nicolas, can you explain the classification again?
I think the collapse to 1d + expand approach is a good one that is quite general. We should only use it when there is no other way to handle the reshape.
I agree. My concern was that in many cases we still need at least a partial flatten then expand, but we can minimize the number of dims we have to flatten. I was expecting that differentiating those cases will become unwieldy so I wanted to see if tensor.reshape was a viable option.
I think the collapse to 1d + expand approach is a good one that is quite general. We should only use it when there is no other way to handle the reshape.
I agree. My concern was that in many cases we still need at least a partial flatten then expand, but we can minimize the number of dims we have to flatten. I was expecting that differentiating those cases will become unwieldy so I wanted to see if tensor.reshape was a viable option.
tensor.reshape is at basically the same abstraction level as the op we are lowering here, so it is definitely easier to lower to. But then somebody else has to do the same analysis that we are doing here to get the expand/collapsing dims. I don't see that happening reliably elsewhere in the ecosystem (I could be wrong), so I would prefer that we "own" that step for now.
tensor.reshape is at basically the same abstraction level as the op we are lowering here, so it is definitely easier to lower to. But then somebody else has to do the same analysis that we are doing here to get the expand/collapsing dims. I don't see that happening reliably elsewhere in the ecosystem (I could be wrong), so I would prefer that we "own" that step for now.
Yep yep, this makes a lot of sense to me. Because this is a common I was more or less wondering whether this should be owned by tensor.reshape and then other pipelines get a benefit if they pass through there, but I'm also fine with us owning it. I think the only difference between aten.view and tensor.reshape is the possibility of inferred dimensions for aten.view.
I'm not actually sure if this PR will turn into the change we want but I'll leave it up for now. I think @JakopinA said he would take a look at this while I'm focused elsewhere.