Refine static shapes for return tensors of conv2d and maxpool2d #502
Conversation
ramiro050
left a comment
ramiro050
left a comment
There was a problem hiding this comment.
Thanks for this! Everything looks good to me, I just have a small style comment below
| } | ||
|
|
||
| // dim_out = | ||
| // floor((dim_in + 2 * padding - dilation * (kernelSize - 1) - 1) / stride) + 1 |
There was a problem hiding this comment.
Nit: I don't think this is necessary since this information can very easily be determined from the body of the function
ramiro050
left a comment
ramiro050
left a comment
There was a problem hiding this comment.
Oh also, given that the size calculations are really complicated, it would be nice to have an e2e test with static shapes to make sure our code agrees with PyTorch
https://github.com/llvm/torch-mlir/blob/main/e2e_testing/torchscript/conv.py
ljfitz
force-pushed
the
static_shapes
branch
from
e1ec961 to
d3f351e
Compare
Thanks @ramiro050 for the review and the suggestion. I've added an e2e test and manually verified that the output of |
ramiro050
left a comment
ramiro050
left a comment
There was a problem hiding this comment.
LGTM! Thanks for adding this!
And thanks for the supportive review. I see that I haven't permission to merge the the pull request. Do you know if merging will happen automatically or if I need to take some additional action? |
Merged! |
2 participants
In RefineTypes.cpp, the return
dtypewas being determined for conv2d and maxpool2d, but the shape was not. This code adds static shape computation then the operator arguments provide the required static details.This may relate to/clash with other possible approaches to #415, where @silvasean talks about revamping shape inference to be more precise, but I don't know what @silvasean had in mind.