[Nested Tensor] view: basic support for ragged_idx != 1 and _unsafe_view #118317

davidberard98 · 2024-01-25T19:29:50Z

Stack from ghstack (oldest at bottom):

-> [Nested Tensor] view: basic support for ragged_idx != 1 and _unsafe_view #118317

Uses case: _unsafe_view is used in aot_autograd to create a view that doesn't register as a view:

pytorch/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py

Lines 470 to 476 in eebe7e1

    
           if CompiledFunction.metadata.num_unsafe_view_outputs > 0: 
        
               for idx in CompiledFunction.metadata.unsafe_view_out_indices: 
        
                   raw_return_idx = num_mutated_runtime_inps + idx 
        
                   o = raw_returns[raw_return_idx] 
        
                   raw_returns[raw_return_idx] = torch.ops.aten._unsafe_view( 
        
                       o, o.shape 
        
                   )

If a transposed nested tensor (i.e. NT with ragged_idx != 1) encounters this code path, it previously would fail for two reasons: 1) because _unsafe_view isn't registered, and 2) because ragged_idx != 1 is not supported. This PR adds support for _unsafe_view (completely reusing the implementation of view; this just registers _unsafe_view as another op using the same implementation). It also adds support for ragged_idx != 1, but only for trivial cases where inp._size == size (the use case used by aot_autograd).

Tests: verify that the result of _unsafe_view doesn't have a _base, and that simple views on transposed NTs work.

Differential Revision: D53096814

…view Uses case: `_unsafe_view` is used in aot_autograd to create a view that doesn't register as a view: https://github.com/pytorch/pytorch/blob/eebe7e1d37f1baa995c694d540cc2fc98884fa18/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py#L470-L476 If a transposed nested tensor (i.e. NT with ragged_idx != 1) encounters this code path, it previously would fail for two reasons: 1) because `_unsafe_view` isn't registered, and 2) because ragged_idx != 1 is not supported. This PR adds support for `_unsafe_view` (completely reusing the implementation of `view`; this just registers `_unsafe_view` as another op using the same implementation). It also adds support for ragged_idx != 1, but only for trivial cases where inp._size == size (the use case used by aot_autograd). Tests: verify that the result of `_unsafe_view` doesn't have a `_base`, and that simple views on transposed NTs work. [ghstack-poisoned]

pytorch-bot · 2024-01-25T19:29:55Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/118317

📄 Preview Python docs built from this PR
📄 Preview C++ docs built from this PR
❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 8428891 with merge base 3cdd4e2 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

davidberard98 · 2024-01-25T19:34:04Z

@davidberard98 has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

…nd _unsafe_view" Uses case: `_unsafe_view` is used in aot_autograd to create a view that doesn't register as a view: https://github.com/pytorch/pytorch/blob/eebe7e1d37f1baa995c694d540cc2fc98884fa18/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py#L470-L476 If a transposed nested tensor (i.e. NT with ragged_idx != 1) encounters this code path, it previously would fail for two reasons: 1) because `_unsafe_view` isn't registered, and 2) because ragged_idx != 1 is not supported. This PR adds support for `_unsafe_view` (completely reusing the implementation of `view`; this just registers `_unsafe_view` as another op using the same implementation). It also adds support for ragged_idx != 1, but only for trivial cases where inp._size == size (the use case used by aot_autograd). Tests: verify that the result of `_unsafe_view` doesn't have a `_base`, and that simple views on transposed NTs work. Differential Revision: [D53096814](https://our.internmc.facebook.com/intern/diff/D53096814) [ghstack-poisoned]

…view Uses case: `_unsafe_view` is used in aot_autograd to create a view that doesn't register as a view: https://github.com/pytorch/pytorch/blob/eebe7e1d37f1baa995c694d540cc2fc98884fa18/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py#L470-L476 If a transposed nested tensor (i.e. NT with ragged_idx != 1) encounters this code path, it previously would fail for two reasons: 1) because `_unsafe_view` isn't registered, and 2) because ragged_idx != 1 is not supported. This PR adds support for `_unsafe_view` (completely reusing the implementation of `view`; this just registers `_unsafe_view` as another op using the same implementation). It also adds support for ragged_idx != 1, but only for trivial cases where inp._size == size (the use case used by aot_autograd). Tests: verify that the result of `_unsafe_view` doesn't have a `_base`, and that simple views on transposed NTs work. ghstack-source-id: 7c15f5b Pull Request resolved: #118317

soulitzer

SGTM

davidberard98 · 2024-01-26T05:18:04Z

@pytorchbot merge -r

pytorchmergebot · 2024-01-26T05:20:56Z

@pytorchbot started a rebase job onto refs/remotes/origin/viable/strict. Check the current status here

…d _unsafe_view" Uses case: `_unsafe_view` is used in aot_autograd to create a view that doesn't register as a view: https://github.com/pytorch/pytorch/blob/eebe7e1d37f1baa995c694d540cc2fc98884fa18/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py#L470-L476 If a transposed nested tensor (i.e. NT with ragged_idx != 1) encounters this code path, it previously would fail for two reasons: 1) because `_unsafe_view` isn't registered, and 2) because ragged_idx != 1 is not supported. This PR adds support for `_unsafe_view` (completely reusing the implementation of `view`; this just registers `_unsafe_view` as another op using the same implementation). It also adds support for ragged_idx != 1, but only for trivial cases where inp._size == size (the use case used by aot_autograd). Tests: verify that the result of `_unsafe_view` doesn't have a `_base`, and that simple views on transposed NTs work. Differential Revision: [D53096814](https://our.internmc.facebook.com/intern/diff/D53096814) [ghstack-poisoned]

pytorchmergebot · 2024-01-26T05:21:09Z

Successfully rebased gh/davidberard98/254/orig onto refs/remotes/origin/viable/strict, please pull locally before adding more changes (for example, via ghstack checkout https://github.com/pytorch/pytorch/pull/118317)

…view Uses case: `_unsafe_view` is used in aot_autograd to create a view that doesn't register as a view: https://github.com/pytorch/pytorch/blob/eebe7e1d37f1baa995c694d540cc2fc98884fa18/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py#L470-L476 If a transposed nested tensor (i.e. NT with ragged_idx != 1) encounters this code path, it previously would fail for two reasons: 1) because `_unsafe_view` isn't registered, and 2) because ragged_idx != 1 is not supported. This PR adds support for `_unsafe_view` (completely reusing the implementation of `view`; this just registers `_unsafe_view` as another op using the same implementation). It also adds support for ragged_idx != 1, but only for trivial cases where inp._size == size (the use case used by aot_autograd). Tests: verify that the result of `_unsafe_view` doesn't have a `_base`, and that simple views on transposed NTs work. ghstack-source-id: 1e0b057 Pull Request resolved: #118317

pytorchmergebot · 2024-01-26T05:22:38Z

Merge failed

Reason: This PR has internal changes and must be landed via Phabricator

Details for Dev Infra team

Raised by workflow job

davidberard98 · 2024-01-26T07:02:09Z

@davidberard98 has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

davidberard98 · 2024-01-26T07:02:58Z

@pytorchbot merge

pytorchmergebot · 2024-01-26T07:05:03Z

Merge failed

Reason: This PR has internal changes and must be landed via Phabricator

Details for Dev Infra team

Raised by workflow job

davidberard98 · 2024-01-26T17:26:57Z

@pytorchbot merge

pytorchmergebot · 2024-01-26T17:29:13Z

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

…iew (pytorch#118317) Uses case: `_unsafe_view` is used in aot_autograd to create a view that doesn't register as a view: https://github.com/pytorch/pytorch/blob/eebe7e1d37f1baa995c694d540cc2fc98884fa18/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py#L470-L476 If a transposed nested tensor (i.e. NT with ragged_idx != 1) encounters this code path, it previously would fail for two reasons: 1) because `_unsafe_view` isn't registered, and 2) because ragged_idx != 1 is not supported. This PR adds support for `_unsafe_view` (completely reusing the implementation of `view`; this just registers `_unsafe_view` as another op using the same implementation). It also adds support for ragged_idx != 1, but only for trivial cases where inp._size == size (the use case used by aot_autograd). Tests: verify that the result of `_unsafe_view` doesn't have a `_base`, and that simple views on transposed NTs work. Differential Revision: [D53096814](https://our.internmc.facebook.com/intern/diff/D53096814) Pull Request resolved: pytorch#118317 Approved by: https://github.com/soulitzer

davidberard98 mentioned this pull request Jan 25, 2024

[NestedTensor] Support ragged_idx != 1 in pointwise ops #118157

Closed

davidberard98 changed the title ~~[Nested Tensor] view: simple support for ragged_idx != 1 and _unsafe_view~~ [Nested Tensor] view: basic support for ragged_idx != 1 and _unsafe_view Jan 25, 2024

davidberard98 added topic: not user facing topic category release notes: jit release notes category labels Jan 25, 2024

davidberard98 requested review from YuqingJ, jbschlosser and soulitzer January 25, 2024 23:55

davidberard98 marked this pull request as ready for review January 25, 2024 23:56

soulitzer approved these changes Jan 26, 2024

View reviewed changes

pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label Jan 26, 2024

pytorchmergebot added the merging label Jan 26, 2024

pytorchmergebot removed the merging label Jan 26, 2024

pytorchmergebot added the merging label Jan 26, 2024

pytorchmergebot removed the merging label Jan 26, 2024

pytorchmergebot added the merging label Jan 26, 2024

pytorchmergebot closed this in 2842d3c Jan 26, 2024

pytorchmergebot added Merged and removed merging labels Jan 26, 2024

facebook-github-bot deleted the gh/davidberard98/254/head branch January 30, 2024 15:22

	if CompiledFunction.metadata.num_unsafe_view_outputs > 0:
	for idx in CompiledFunction.metadata.unsafe_view_out_indices:
	raw_return_idx = num_mutated_runtime_inps + idx
	o = raw_returns[raw_return_idx]
	raw_returns[raw_return_idx] = torch.ops.aten._unsafe_view(
	o, o.shape
	)

[Nested Tensor] view: basic support for ragged_idx != 1 and _unsafe_view #118317

[Nested Tensor] view: basic support for ragged_idx != 1 and _unsafe_view #118317

Uh oh!

Conversation

davidberard98 commented Jan 25, 2024 • edited by pytorchmergebot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Jan 25, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/118317

✅ No Failures

Uh oh!

davidberard98 commented Jan 25, 2024

Uh oh!

soulitzer left a comment

Choose a reason for hiding this comment

Uh oh!

davidberard98 commented Jan 26, 2024

Uh oh!

pytorchmergebot commented Jan 26, 2024

Uh oh!

pytorchmergebot commented Jan 26, 2024

Uh oh!

pytorchmergebot commented Jan 26, 2024

Merge failed

Uh oh!

davidberard98 commented Jan 26, 2024

Uh oh!

davidberard98 commented Jan 26, 2024

Uh oh!

pytorchmergebot commented Jan 26, 2024

Merge failed

Uh oh!

davidberard98 commented Jan 26, 2024

Uh oh!

pytorchmergebot commented Jan 26, 2024

Merge started

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

davidberard98 commented Jan 25, 2024 •

edited by pytorchmergebot

Loading

pytorch-bot bot commented Jan 25, 2024 •

edited

Loading