Skip to content

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

Merged
zjgarvey merged 6 commits into
llvm:mainfrom
zahidwx:feature/ReplicationPad1D
Jun 15, 2025
Merged

[TorchToLinalg] Support AtenReplicationPad1d with lowering to linalg backend#4217
zjgarvey merged 6 commits into
llvm:mainfrom
zahidwx:feature/ReplicationPad1D

Conversation

@zahidwx

@zahidwx zahidwx commented Jun 1, 2025

Copy link
Copy Markdown
Contributor

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

zahidwx commented Jun 10, 2025

Copy link
Copy Markdown
Contributor Author

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

@zjgarvey zjgarvey left a comment

Copy link
Copy Markdown
Collaborator

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.

Comment thread lib/Conversion/TorchToLinalg/TensorConstructors.cpp Outdated
Comment thread lib/Conversion/TorchToLinalg/TensorConstructors.cpp
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
Comment thread lib/Conversion/TorchToLinalg/TensorConstructors.cpp Outdated
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
Copy Markdown
Collaborator

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

@zahidwx

zahidwx commented Jun 15, 2025

Copy link
Copy Markdown
Contributor Author

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