fix: infer output shape directly in ConvertAtenUnflattenIntOp #4325
Conversation
| expand = | ||
| rewriter.create<tensor::CastOp>(loc, expandTy, expand).getResult(); | ||
| } | ||
| } else { |
There was a problem hiding this comment.
Is this path still needed?
There was a problem hiding this comment.
Yes, I think so. I tried without, doesn't seem to be possible. Unless you had a specific patch in mind?
There was a problem hiding this comment.
I don't see when it would need to create a tensor.reshape. Could you add a test to unflatten.mlir for when the tensor.reshape is needed?
There was a problem hiding this comment.
To be honest I'm not sure either outside of the explicit check, for which I've added a case for now. But otherwise I'm not sure if it's possible to simplify the logic. Sometime last week I tested trying to combine the two and it didn't work.
Yes, will do. |
Added some basic tests - I think there are some other cases that we need to add but let me know if that's what you had in mind |
There was a problem hiding this comment.
Thanks for the PR.
I'm a bit out-of-context on the issue, but this seems like it is not a problem with the lowering of unflatten, but rather an issue with a decomposition generating this op without a fully specified shape. There are shape inference passes which should automatically resolve the output shape if it can be inferred statically, but we collectively decided as a group to not require results of decompose-complex-ops to undergo an addition shape refinement application (these are a bit time-expensive).
I'll comment some more details on the yolo issue you linked, but I don't think the present changes are going to be the best approach to resolving the issue.
|
|
||
| ```shell | ||
| cmake --build build --target check-torch-mlir-python | ||
| cmake --build build --target check-torch_mlir-python |
There was a problem hiding this comment.
I don't really understand this change. If the old script is broken (which would be very odd), this should definitely be made into a separate PR.
There was a problem hiding this comment.
I just noticed that trying cmake --build build --target check-torch-mlir-python would ask "Did you mean check-torch_mlir-python" so I made that change. But I will make this a separate PR.
|
@raayandhar I posted some more info here: iree-org/iree#22022 (comment) It's likely that you should submit a quick fix for 1., and it would be great to pick up 2., but this might be a bit of a difficult task to handle properly. |
Thanks for the help and review Zach, will get a fix out for 1 today. I'll have to understand 2. better first before trying to tackle it. |
We are facing an error that is related to a specific unflatten case - when we have dimensions that need to be inferred, we cannot handle them at the moment. For example, input shape [128] and sizes [1, 1, 1, -1], the -1 is to infer the dimension in that slot, and we expect to get [1, 1, 1, 128]. We calculate this directly and use this to build our
ExpandShapeOp.Some more discussion:
aten.as_stridedcauses assertion failure in MLIR #4315 and this issue is related to [MLIR][TORCH] Add E2E support for aten.as_strided op #4269