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.

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

Lower AtenFull op #5781

Merged
merged 3 commits into from
Nov 14, 2023
Merged

Lower AtenFull op #5781

merged 3 commits into from
Nov 14, 2023

Conversation

danielvegamyhre
Copy link
Collaborator

@danielvegamyhre danielvegamyhre commented Nov 8, 2023

Fixes #5780

This is my first attempt at op lowering, and I have a couple questions:

  1. Do we want to fallback to the pytorch/aten CPU version of this op if the layout and/or pin_memory values are not the default? I saw this pattern in other places in in aten_xla_type.cpp, and the tensor_methods::full function in XLA doesn't support these parameters, so it seemed like the right thing to do, but I'm not sure.

  2. I notice full was not listed in xla_native_functions.yaml (as expected) but there IS an existing test for it in test_aten_xla_tensor_1.cpp called TestFull, which seems to verify that we use xla::empty followed by xla::fill_ to create a tensor with the fill value. So I just changed this test to check that we are instead using the new full method implemented in this PR. I confirmed the test is passing with these changes, and failing without the changes. Is this the correct thing to do here?

Copy link
Collaborator

@wonjoolee95 wonjoolee95 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 working on this, @danielvegamyhre!

  1. Do we want to fallback to the pytorch/aten CPU version of this op if the layout and/or pin_memory values are not the default? I saw this pattern in other places in in aten_xla_type.cpp, and the tensor_methods::full function in XLA doesn't support these parameters, so it seemed like the right thing to do, but I'm not sure.

Yes, this is the correct behavior. We can fallback to CPU for now for full like you did in this PR.

  1. I notice full was not listed in xla_native_functions.yaml (as expected) but there IS an existing test for it in test_aten_xla_tensor_1.cpp called TestFull, which seems to verify that we use xla::empty followed by xla::fill_ to create a tensor with the fill value. So I just changed this test to check that we are instead using the new full method implemented in this PR. I confirmed the test is passing with these changes, and failing without the changes. Is this the correct thing to do here?

Good find -- so most likely what's happening here is that when full is not lowered, it's decomposed empty and fill. But since this PR lowers full, what you did here -- removing the counter checks for empty and fill -- is correct, thanks!

LGTM.

@wonjoolee95 wonjoolee95 merged commit b3ed82c into pytorch:master Nov 14, 2023
17 checks passed
@miladm
Copy link
Collaborator

miladm commented Nov 14, 2023

Awesome! Thank you @danielvegamyhre

mbzomowski pushed a commit to mbzomowski-test-org/xla that referenced this pull request Nov 16, 2023
* lower full

* update test for full op

* formatting
zpcore pushed a commit that referenced this pull request Nov 21, 2023
* lower full

* update test for full op

* formatting
lsy323 pushed a commit to lsy323/xla that referenced this pull request Nov 28, 2023
* lower full

* update test for full op

* formatting
chunnienc pushed a commit to chunnienc/xla that referenced this pull request Dec 14, 2023
* lower full

* update test for full op

* formatting
golechwierowicz pushed a commit that referenced this pull request Jan 12, 2024
* lower full

* update test for full op

* formatting
bhavya01 pushed a commit that referenced this pull request Apr 22, 2024
* lower full

* update test for full op

* formatting
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.

Op lowering for AtenFull
3 participants