Skip to content

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

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. #546

wants to merge 2 commits into from

Conversation

@gprateek93
Copy link
Collaborator

@gprateek93 gprateek93 commented Jan 28, 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 gprateek93 force-pushed the prateek/native-layer-norm-backward branch from 83a4094 to 8dc4dc4 Compare January 28, 2022 11:12
@gprateek93 gprateek93 mentioned this pull request Jan 28, 2022
Closed
@gprateek93 gprateek93 force-pushed the prateek/native-layer-norm-backward branch 2 times, most recently from c9fee2e to aaaef5d Compare January 28, 2022 12:40
pashu123
pashu123 requested changes Jan 29, 2022
View reviewed changes
cathyzhyi
cathyzhyi requested changes Jan 31, 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.

Left some comments. Haven't done reviewing yet.

cathyzhyi
cathyzhyi reviewed Jan 31, 2022
View reviewed changes
@gprateek93 gprateek93 force-pushed the prateek/native-layer-norm-backward branch 3 times, most recently from 624de1e to 5a9f66f Compare February 2, 2022 15:47
cathyzhyi
cathyzhyi reviewed Feb 3, 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.

Still haven't done reviewing but I am wondering if you've considered implementing this in DecomposeComplexOps?

cathyzhyi
cathyzhyi reviewed Feb 3, 2022
View reviewed changes
@gprateek93
Copy link
Collaborator Author

gprateek93 commented Feb 3, 2022 •
edited
Loading

Still haven't done reviewing but I am wondering if you've considered implementing this in DecomposeComplexOps?

I was following the previous code of batchNorm and layerNorm. But you are right, AFAIK these ops can be done in the DecomposeComplexOps pass. I can add a patch in this PR to make such change.

@gprateek93 gprateek93 force-pushed the prateek/native-layer-norm-backward branch 2 times, most recently from f9dbf33 to 04bfc4c Compare February 9, 2022 09:11
@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 force-pushed the prateek/native-layer-norm-backward branch from 04bfc4c to d2cd480 Compare February 9, 2022 09:19
@gprateek93
Copy link
Collaborator Author

gprateek93 commented Feb 9, 2022

Still haven't done reviewing but I am wondering if you've considered implementing this in DecomposeComplexOps?

Hi @cathyzhyi I tried to use the DecomposeComplexOps, but the preprocessing(Checking of requirements) of both the lowerings viz aten.native_layer_norm and its backward op require some std dialect ops like AssertOp.

@cathyzhyi
Copy link
Contributor

cathyzhyi commented Feb 9, 2022

Still haven't done reviewing but I am wondering if you've considered implementing this in DecomposeComplexOps?

Hi @cathyzhyi I tried to use the DecomposeComplexOps, but the preprocessing(Checking of requirements) of both the lowerings viz aten.native_layer_norm and its backward op require some std dialect ops like AssertOp.

Yea, @Shukla-Gaurav has a PR for adding assert into Torch dialect #564. We can use that once it's merged.

@gprateek93
Copy link
Collaborator Author

gprateek93 commented Feb 9, 2022

Opened the PR at (#570)

@gprateek93 gprateek93 closed this Feb 9, 2022
@gprateek93 gprateek93 deleted the prateek/native-layer-norm-backward branch February 9, 2022 16:02
qedawkins pushed a commit to nod-ai/torch-mlir that referenced this pull request Oct 3, 2022
@chentong319
Clean attribute promote (llvm#546)
6da941f 
* change dialect

Signed-off-by: chentong <chentong@us.ibm.com>

* fix pad

Signed-off-by: chentong <chentong@us.ibm.com>

* Reshape

Signed-off-by: chentong <chentong@us.ibm.com>

* format

Signed-off-by: chentong <chentong@us.ibm.com>

* adjust test

Signed-off-by: chentong <chentong@us.ibm.com>

* docs

Signed-off-by: chentong <chentong@us.ibm.com>

* remove the code

Signed-off-by: chentong <chentong@us.ibm.com>

* clean gen_onnx_mlir

Signed-off-by: chentong <chentong@us.ibm.com>

* format

Signed-off-by: chentong <chentong@us.ibm.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@cathyzhyi cathyzhyi cathyzhyi requested changes

@pashu123 pashu123 pashu123 requested changes

@ramiro050 ramiro050 Awaiting requested review from ramiro050

@Shukla-Gaurav Shukla-Gaurav Awaiting requested review from Shukla-Gaurav

@silvasean silvasean Awaiting requested review from silvasean

@vivekkhandelwal1 vivekkhandelwal1 Awaiting requested review from vivekkhandelwal1

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

@gprateek93 @cathyzhyi @pashu123