Skip to content

[TorchToLinalg] Support AtenReplicationPad1d with lowering to linalg backend #4217

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

Merged
merged 6 commits into from
Jun 15, 2025

Conversation

zahidwx
Copy link
Contributor

@zahidwx zahidwx commented Jun 1, 2025

This PR takes care of #4216.

  • Add AtenReplicationPad1d support in torch dialect.
  • Lower to linalg backend:
    • tensor.extract_slice ops for left & right slices
    • tensor.concat ops to expand each slice into right shape and join all slices in order
  • Update AtenPadOp decomposition to make use of 1d replication pad instead of using 2d variant.
  • Add tests & update failing cases in xfail sets.

@zahidwx zahidwx force-pushed the feature/ReplicationPad1D branch from 47089ef to e2e701b Compare June 1, 2025 20:11
@zahidwx zahidwx marked this pull request as ready for review June 1, 2025 20:50
@zahidwx
Copy link
Contributor Author

zahidwx commented Jun 10, 2025

Would appreciate a review - @vivekkhandelwal1 | @stellaraccident | @penguin-wwy

Copy link
Collaborator

@zjgarvey zjgarvey left a comment

Choose a reason for hiding this comment

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

Thanks for the PR!

I'm not entirely sure about the original motivation for the 2d lowering, but 3 different concat ops seems like an expensive way to perform padding.

I'm not going to block this either way. The PR looks good in general, and it would make more sense to add a lowering consistent with the 2d op first, and follow up with a separate patch to improve the implementation if desired.

zahidwx added 5 commits June 14, 2025 12:51
Signed-off-by: Zahid Wakeel <zahid.wakeel@multicorewareinc.com>
Signed-off-by: Zahid Wakeel <zahid.wakeel@multicorewareinc.com>
Signed-off-by: Zahid Wakeel <zahid.wakeel@multicorewareinc.com>
Signed-off-by: Zahid Wakeel <zahid.wakeel@multicorewareinc.com>
Signed-off-by: Zahid Wakeel <zahid.wakeel@multicorewareinc.com>
@zahidwx zahidwx force-pushed the feature/ReplicationPad1D branch from c8af8a9 to 4a9bc26 Compare June 14, 2025 07:23
Signed-off-by: Zahid Wakeel <zahid.wakeel@multicorewareinc.com>
@zahidwx zahidwx force-pushed the feature/ReplicationPad1D branch from 04aa077 to 691eb58 Compare June 14, 2025 18:54
@zjgarvey
Copy link
Collaborator

My approval still stands, do you want me to merge?

@zahidwx
Copy link
Contributor Author

zahidwx commented Jun 15, 2025

My approval still stands, do you want me to merge?

Yes please @zjgarvey

@zjgarvey zjgarvey merged commit 3850332 into llvm:main Jun 15, 2025
3 checks passed
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.

2 participants