Skip to content

[TORCH][MLIR] Add E2E support for aten.native_layer_norm_backward op. #570

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
Closed
wants to merge 2 commits into from
Closed

[TORCH][MLIR] Add E2E support for aten.native_layer_norm_backward op. #570

wants to merge 2 commits into from

Conversation

@gprateek93
Copy link
Collaborator

@gprateek93 gprateek93 commented Feb 9, 2022

This PR adds two commits:

  1. adds support for aten.native_layer_norm_backward operation. It also adds support for matching constant bools stored in a boolean list.
  2. fixes aten.native_layer_norm. Previously this operation was not calculating correct shapes for mean and inverted STD. This has been corrected in this commit. Some new helper functions are added to calculate the inverted STD and to broadcast a given input with the help of a broadcast mask.

Signed-Off-By: Prateek Gupta prateek@nod-labs.com

@gprateek93
[TORCH][MLIR] Fix the lowering of aten.native_layer_norm op.
28a7c7d 
This commit fixs the lowering of `aten.native_layer_norm` operation.
Previously this operation was not calculating correct shapes for
mean and inverted STD. This has been corrected in this commit.
Some new helper functions are added to calculate the inverted STD
and to broadcast a given input with the help of a broadcast mask.

Signed-Off-By: Prateek Gupta <prateek@nod-labs.com>
@gprateek93
[TORCH][MLIR] Add E2E support for aten.native_layer_norm_backward op.
d2cd480 
This commit adds support for `aten.native_layer_norm_backward` operation.
It also adds support for matching constant bools stored in a boolean
list.

Signed-Off-By: Prateek Gupta <prateek@nod-labs.com>
@gprateek93 gprateek93 mentioned this pull request Feb 9, 2022
Closed
@gprateek93 gprateek93 mentioned this pull request Feb 9, 2022
Closed
cathyzhyi
cathyzhyi reviewed Feb 9, 2022
View reviewed changes
Copy link
Contributor

@cathyzhyi cathyzhyi left a comment

Choose a reason for hiding this comment

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

I would recommend implementing this in DecomposeComplexOps if possible. Math calculation in DecomposeComplexOps are easier to implement and more readable and maintainable.

@gprateek93
Copy link
Collaborator Author

gprateek93 commented Mar 10, 2022

@cathyzhyi I have added a PR for this decomposition in functorch: pytorch/functorch#525
So I guess having the decomposition here also will duplicate it.

@cathyzhyi
Copy link
Contributor

cathyzhyi commented Mar 15, 2022

@cathyzhyi I have added a PR for this decomposition in functorch: pytorch/functorch#525 So I guess having the decomposition here also will duplicate it.

@gprateek93 After the decomposition is merged can you add an e2e test to make sure it would make aten::native_layer_norm_backward work?

@cathyzhyi cathyzhyi mentioned this pull request Mar 15, 2022
Closed
@gprateek93 gprateek93 marked this pull request as draft March 29, 2022 16:57
@gprateek93
Copy link
Collaborator Author

gprateek93 commented Mar 30, 2022 •
edited
Loading

The functorch PR for aten.native_layer_norm_backward is merged. Waiting for the integration of latest functorch in torch-mlir. Once done, we can safely close this PR.

@henrytwo
Copy link
Member

henrytwo commented May 18, 2022 •
edited
Loading

Hey are there any updates on this PR? We seem to have support for native_layer_norm but it would be great to also have native_layer_norm_backward working too.

Ideally we'd like to have native_layer_norm_backward stay as a high level op, rather than have it be decomposed.

cc: @antoniojkim @ke1337

@henrytwo henrytwo mentioned this pull request May 30, 2022
Closed
@henrytwo
Copy link
Member

henrytwo commented May 30, 2022 •
edited
Loading

FYI: I'm working on a new PR that will borrow some code from this one (with credits and reference), since I need support for aten.native_layer_norm_backward soon. #888

@henrytwo henrytwo mentioned this pull request May 30, 2022
Merged
qedawkins pushed a commit to nod-ai/torch-mlir that referenced this pull request Oct 3, 2022
@dlwh @AlexandreEichenberger @caoimhinuibrian
fix typo (llvm#570)
1a165ea 
(trivial drive-by fix)

Co-authored-by: Alexandre Eichenberger <alexe@us.ibm.com>
Co-authored-by: Kevin O'Brien <caomhin@us.ibm.com>
@xgupta
Copy link
Contributor

xgupta commented Nov 26, 2022

This PR can be closed, seems #888 is merged to add aten.native_layer_norm_backward op.

@pashu123 pashu123 closed this Nov 26, 2022
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@cathyzhyi cathyzhyi cathyzhyi left review comments

@pashu123 pashu123 Awaiting requested review from pashu123

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

5 participants

@gprateek93 @cathyzhyi @henrytwo @xgupta @pashu123