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.

Sign up for GitHub

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

Jump to bottom

[Torch Dialect] Support Einsum Op #2230

Merged
merged 28 commits into from
Dec 10, 2023
Merged

[Torch Dialect] Support Einsum Op #2230

merged 28 commits into from
Dec 10, 2023

Conversation

JianzheXiao
Copy link
Contributor

@JianzheXiao JianzheXiao commented Jun 12, 2023 •
edited
Loading

As title, support torch.aten.einsum op

Right now only support Static Shape, because of the known issue, the fixed solution is here: #2154

Co-authored-by: Jiawei Wu wujiawei.aml@bytedance.com

ramiro050
ramiro050 requested changes Jun 13, 2023
View reviewed changes
Copy link
Collaborator

@ramiro050 ramiro050 left a comment

Choose a reason for hiding this comment

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

This is a really complex decomposition. Good job getting it working! Initial set of comments

lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp Outdated Show resolved Hide resolved
lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp Outdated Show resolved Hide resolved
lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp Outdated Show resolved Hide resolved
lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp Outdated Show resolved Hide resolved
lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp Outdated Show resolved Hide resolved
lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp Show resolved Hide resolved
lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp Outdated Show resolved Hide resolved
lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp Outdated Show resolved Hide resolved
lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp Outdated Show resolved Hide resolved
lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp Outdated Show resolved Hide resolved
@JianzheXiao
Copy link
Contributor Author

@ramiro050 I just refactor the code, could you please review it again?

ramiro050
ramiro050 requested changes Jul 17, 2023
View reviewed changes
Copy link
Collaborator

@ramiro050 ramiro050 left a comment

Choose a reason for hiding this comment

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

A few more comments

lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp Outdated Show resolved Hide resolved
lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp Show resolved Hide resolved
lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp Outdated Show resolved Hide resolved
lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp Outdated Show resolved Hide resolved
lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp Outdated Show resolved Hide resolved
lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp Outdated Show resolved Hide resolved
python/torch_mlir_e2e_test/test_suite/reshape_like.py Outdated Show resolved Hide resolved
lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp Outdated Show resolved Hide resolved
lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp Outdated Show resolved Hide resolved
lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp Outdated Show resolved Hide resolved
@JianzheXiao
Copy link
Contributor Author

@ramiro050 Could you please review this pr again? Thank you!

Vremold
Vremold reviewed Aug 9, 2023
View reviewed changes
Copy link
Collaborator

@Vremold Vremold left a comment •
edited
Loading

Choose a reason for hiding this comment

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

A few comments plus a question. Can we make sure that all einsum ops like matmul or Hadamard product pass this decomposition pattern, e.g., an einsum op with equation "abc,abc->"?

lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp Outdated Show resolved Hide resolved
lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp Outdated Show resolved Hide resolved
lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp Outdated Show resolved Hide resolved
lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp Outdated Show resolved Hide resolved
lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp Outdated Show resolved Hide resolved
lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp Show resolved Hide resolved
@JianzheXiao
Copy link
Contributor Author

A few comments plus a question. Can we make sure that all einsum ops like matmul or Hadamard product pass this decomposition pattern, e.g., an einsum op with equation "abc,abc->"?

matmul should be okay but right now the algorithm are not able to deal with the case like abc, abc-> abc

ramiro050
ramiro050 requested changes Nov 15, 2023
View reviewed changes
Copy link
Collaborator

@ramiro050 ramiro050 left a comment

Choose a reason for hiding this comment

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

Looks mostly good to me, just one question

lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp Outdated Show resolved Hide resolved
ramiro050
ramiro050 approved these changes Dec 8, 2023
View reviewed changes
Copy link
Collaborator

@ramiro050 ramiro050 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 hard work on this op 🙂

@qingyunqu qingyunqu merged commit 96fcde4 into llvm:main Dec 10, 2023
5 checks passed
@JianzheXiao JianzheXiao deleted the add_support_einsum branch December 12, 2023 00:54
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@Vremold Vremold Vremold left review comments

@qingyunqu qingyunqu qingyunqu approved these changes

@ramiro050 ramiro050 ramiro050 approved these changes

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
4 participants
@JianzheXiao @ramiro050 @qingyunqu @Vremold