Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Im2col] Allow multiple batch, M, and K dimensions on im2col result #18593

Merged
merged 7 commits into from
Oct 9, 2024

Conversation

Max191
Copy link
Contributor

@Max191 Max191 commented Sep 24, 2024

This PR adds support for multiple M and K dimensions in the result of the im2col op. New metadata is added for correctly tracking the offsets into the M and K dimensions along the multiple dimensions. New m_basis and k_basis fields are added to the op, which represent a basis for linearizing the m_offset and k_offset fields.

The motivation for doing this is that flattening the M dimension can create an expand_shape op consumer of the resulting matmul. This can cause issues with fusion and distribution, so it is useful to be able to keep the multiple M dimensions intact. This PR does not change any behavior of Conv2DToIm2col pass, which will be done in a later PR.

@Max191 Max191 force-pushed the expanded-im2col-op branch 4 times, most recently from 83997e3 to 3906d86 Compare September 25, 2024 16:36
@Max191 Max191 marked this pull request as ready for review September 25, 2024 16:36
Copy link
Contributor

@hanhanW hanhanW left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

First round of review comments, please take a look. I still need to ramp up context for a deeper review.

@Max191 Max191 requested a review from qedawkins October 8, 2024 17:09
Copy link
Contributor

@qedawkins qedawkins left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ok this looks fine to me but I think we should keep an eye on this op and evaluate what it is giving us in the long run. It is starting to track a lot of state that might be better to model as a generic gather.

Copy link
Contributor

@hanhanW hanhanW left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

just one nit in verifier.

Signed-off-by: Max Dawkins <max.dawkins@gmail.com>
Signed-off-by: Max Dawkins <max.dawkins@gmail.com>
Signed-off-by: Max Dawkins <max.dawkins@gmail.com>
Signed-off-by: Max Dawkins <max.dawkins@gmail.com>
Signed-off-by: Max Dawkins <max.dawkins@gmail.com>
Signed-off-by: Max Dawkins <max.dawkins@gmail.com>
Signed-off-by: Max Dawkins <max.dawkins@gmail.com>
@Max191 Max191 merged commit 598a60e into iree-org:main Oct 9, 2024
35 checks passed
Max191 added a commit that referenced this pull request Oct 14, 2024
This is following up on #18593,
which enabled support for expanded result shapes in the im2col op. This
PR changes the rewrite patterns for conv into IGEMM to use the expanded
H and W dimensions, rather than inserting collapse_shape and
expand_shape ops.

---------

Signed-off-by: Max Dawkins <max.dawkins@gmail.com>
@Max191 Max191 deleted the expanded-im2col-op branch October 25, 2024 14:19
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants