functionalization bugfix: using owning type when unwrapping tensors by bdhirsh · Pull Request #76125 · pytorch/pytorch

bdhirsh · 2022-04-20T18:57:26Z

Addresses the comment here: #73441 (comment)

This PR should also fix a bug with at::cat.out in functionalization, so I added a test for it.

When the functionalization pass unwraps tensor arguments, we need to make sure the thing that we unwrap into has an owning type - we can't blindly base the type off whatever types come from the operator's schema.

This is basically just a problem for at::TensorList (non-owning) - we need to use an owning std::vector<Tensor> instead.

@ezyang I wasn't sure how general of a fix to make. translate() already knows a bit about "expensive conversions" from non-owning to owning types. I couldn't think of a way to re-use that code without making some larger changes (like teaching CType's about whether they're owning types), so I added a simple fix inside the functionalization codegen. Feedback totally welcome.

Here's part of the old cat kernel:

    at::Tensor & cat_out_out(c10::DispatchKeySet dispatchKeySet, at::TensorList tensors, int64_t dim, at::Tensor & out) {
      at::TensorList tensors_;
      if (at::functionalization::impl::isFunctionalTensor(tensors)) {
        at::functionalization::impl::sync(tensors);
        tensors_ = at::functionalization::impl::from_functional_tensor(tensors);
      } else {
        tensors_ = tensors;
      }
      ....
    }

And here's the new one (the temporary it creates is a vector)

    at::Tensor & cat_out_out(c10::DispatchKeySet dispatchKeySet, at::TensorList tensors, int64_t dim, at::Tensor & out) {
      ::std::vector<at::Tensor> tensors_;
      if (at::functionalization::impl::isFunctionalTensor(tensors)) {
        at::functionalization::impl::sync(tensors);
        tensors_ = at::functionalization::impl::from_functional_tensor(tensors);
      } else {
        tensors_ = tensors;
      }
      ...
    }

Stack from ghstack:

integrate functionalization <> LTC torchscript backend #75527 [prototype] integrate functionalization <> LTC torchscript backend
generate out= and functional variants of NativeFunctions, get functionalization to work for all mutable ops #76320 attempt to functionalize ops with mutable positional-only args
add native view_copy.out ops, teach codegen about tensorlist out= #76126 add native view_copy.out ops, teach codegen about tensorlist out=
fix torch.tensor for functionalization #76319 [poc] try to fix torch.tensor for functionalization
fix nested grad(functionalize(f)) transforms #76318 fix nested grad(functionalize(f)) transforms
functionalization bugfix: using owning type when unwrapping tensors #76125 functionalization bugfix: using owning type when unwrapping tensors

[ghstack-poisoned]

facebook-github-bot · 2022-04-20T18:57:31Z

🔗 Helpful links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/76125
📄 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 our office hours

💊 CI failures summary and remediations

As of commit ebb412f (more details on the Dr. CI page):

Expand to see more

💚 💚 Looks good so far! There are no failures yet. 💚 💚

This comment was automatically generated by Dr. CI (expand for details).

Please report bugs/suggestions to the (internal) Dr. CI Users group.

Click here to manually regenerate this comment.

bdhirsh · 2022-04-20T18:58:41Z

aten/src/ATen/FunctionalTensorWrapper.cpp

+  for (const auto i : c10::irange(functional_tensor.size())) {
+    commit_update(functional_tensor[i]);
+  }
+}


These aren't strictly necessary for this PR, but I need them later (accidentally bundled them in this commit)

ezyang

the one off seems fine. keep an eye on it in the future

@ezyang

…g tensors" Addresses the comment here: #73441 (comment) When the functionalization pass unwraps tensor arguments, we need to make sure the thing that we unwrap into has an owning type - we can't blindly base the type off whatever types come from the operator's schema. This is basically just a problem for `at::TensorList` (non-owning) - we need to use an owning `std::vector<Tensor>` instead. @ezyang I wasn't sure how general of a fix to make. `translate()` already knows a bit about "expensive conversions" from non-owning to owning types. I couldn't think of a way to re-use that code without making some larger changes (like teaching `CType`'s about whether they're owning types), so I added a simple fix inside the functionalization codegen. Feedback totally welcome. Here's part of the old `cat` kernel: ``` at::Tensor & cat_out_out(c10::DispatchKeySet dispatchKeySet, at::TensorList tensors, int64_t dim, at::Tensor & out) { at::TensorList tensors_; if (at::functionalization::impl::isFunctionalTensor(tensors)) { at::functionalization::impl::sync(tensors); tensors_ = at::functionalization::impl::from_functional_tensor(tensors); } else { tensors_ = tensors; } .... } ``` And here's the new one (the temporary it creates is a vector) ``` at::Tensor & cat_out_out(c10::DispatchKeySet dispatchKeySet, at::TensorList tensors, int64_t dim, at::Tensor & out) { ::std::vector<at::Tensor> tensors_; if (at::functionalization::impl::isFunctionalTensor(tensors)) { at::functionalization::impl::sync(tensors); tensors_ = at::functionalization::impl::from_functional_tensor(tensors); } else { tensors_ = tensors; } ... } ``` [ghstack-poisoned]

@ezyang

…g tensors" Addresses the comment here: #73441 (comment) When the functionalization pass unwraps tensor arguments, we need to make sure the thing that we unwrap into has an owning type - we can't blindly base the type off whatever types come from the operator's schema. This is basically just a problem for `at::TensorList` (non-owning) - we need to use an owning `std::vector<Tensor>` instead. @ezyang I wasn't sure how general of a fix to make. `translate()` already knows a bit about "expensive conversions" from non-owning to owning types. I couldn't think of a way to re-use that code without making some larger changes (like teaching `CType`'s about whether they're owning types), so I added a simple fix inside the functionalization codegen. Feedback totally welcome. Here's part of the old `cat` kernel: ``` at::Tensor & cat_out_out(c10::DispatchKeySet dispatchKeySet, at::TensorList tensors, int64_t dim, at::Tensor & out) { at::TensorList tensors_; if (at::functionalization::impl::isFunctionalTensor(tensors)) { at::functionalization::impl::sync(tensors); tensors_ = at::functionalization::impl::from_functional_tensor(tensors); } else { tensors_ = tensors; } .... } ``` And here's the new one (the temporary it creates is a vector) ``` at::Tensor & cat_out_out(c10::DispatchKeySet dispatchKeySet, at::TensorList tensors, int64_t dim, at::Tensor & out) { ::std::vector<at::Tensor> tensors_; if (at::functionalization::impl::isFunctionalTensor(tensors)) { at::functionalization::impl::sync(tensors); tensors_ = at::functionalization::impl::from_functional_tensor(tensors); } else { tensors_ = tensors; } ... } ``` [ghstack-poisoned]

@ezyang

…g tensors" Addresses the comment here: #73441 (comment) When the functionalization pass unwraps tensor arguments, we need to make sure the thing that we unwrap into has an owning type - we can't blindly base the type off whatever types come from the operator's schema. This is basically just a problem for `at::TensorList` (non-owning) - we need to use an owning `std::vector<Tensor>` instead. @ezyang I wasn't sure how general of a fix to make. `translate()` already knows a bit about "expensive conversions" from non-owning to owning types. I couldn't think of a way to re-use that code without making some larger changes (like teaching `CType`'s about whether they're owning types), so I added a simple fix inside the functionalization codegen. Feedback totally welcome. Here's part of the old `cat` kernel: ``` at::Tensor & cat_out_out(c10::DispatchKeySet dispatchKeySet, at::TensorList tensors, int64_t dim, at::Tensor & out) { at::TensorList tensors_; if (at::functionalization::impl::isFunctionalTensor(tensors)) { at::functionalization::impl::sync(tensors); tensors_ = at::functionalization::impl::from_functional_tensor(tensors); } else { tensors_ = tensors; } .... } ``` And here's the new one (the temporary it creates is a vector) ``` at::Tensor & cat_out_out(c10::DispatchKeySet dispatchKeySet, at::TensorList tensors, int64_t dim, at::Tensor & out) { ::std::vector<at::Tensor> tensors_; if (at::functionalization::impl::isFunctionalTensor(tensors)) { at::functionalization::impl::sync(tensors); tensors_ = at::functionalization::impl::from_functional_tensor(tensors); } else { tensors_ = tensors; } ... } ``` [ghstack-poisoned]

@ezyang

…g tensors" Addresses the comment here: #73441 (comment) When the functionalization pass unwraps tensor arguments, we need to make sure the thing that we unwrap into has an owning type - we can't blindly base the type off whatever types come from the operator's schema. This is basically just a problem for `at::TensorList` (non-owning) - we need to use an owning `std::vector<Tensor>` instead. @ezyang I wasn't sure how general of a fix to make. `translate()` already knows a bit about "expensive conversions" from non-owning to owning types. I couldn't think of a way to re-use that code without making some larger changes (like teaching `CType`'s about whether they're owning types), so I added a simple fix inside the functionalization codegen. Feedback totally welcome. Here's part of the old `cat` kernel: ``` at::Tensor & cat_out_out(c10::DispatchKeySet dispatchKeySet, at::TensorList tensors, int64_t dim, at::Tensor & out) { at::TensorList tensors_; if (at::functionalization::impl::isFunctionalTensor(tensors)) { at::functionalization::impl::sync(tensors); tensors_ = at::functionalization::impl::from_functional_tensor(tensors); } else { tensors_ = tensors; } .... } ``` And here's the new one (the temporary it creates is a vector) ``` at::Tensor & cat_out_out(c10::DispatchKeySet dispatchKeySet, at::TensorList tensors, int64_t dim, at::Tensor & out) { ::std::vector<at::Tensor> tensors_; if (at::functionalization::impl::isFunctionalTensor(tensors)) { at::functionalization::impl::sync(tensors); tensors_ = at::functionalization::impl::from_functional_tensor(tensors); } else { tensors_ = tensors; } ... } ``` [ghstack-poisoned]

@ezyang

…g tensors" Addresses the comment here: #73441 (comment) This PR should also fix a bug with `at::cat.out` in functionalization, so I added a test for it. When the functionalization pass unwraps tensor arguments, we need to make sure the thing that we unwrap into has an owning type - we can't blindly base the type off whatever types come from the operator's schema. This is basically just a problem for `at::TensorList` (non-owning) - we need to use an owning `std::vector<Tensor>` instead. @ezyang I wasn't sure how general of a fix to make. `translate()` already knows a bit about "expensive conversions" from non-owning to owning types. I couldn't think of a way to re-use that code without making some larger changes (like teaching `CType`'s about whether they're owning types), so I added a simple fix inside the functionalization codegen. Feedback totally welcome. Here's part of the old `cat` kernel: ``` at::Tensor & cat_out_out(c10::DispatchKeySet dispatchKeySet, at::TensorList tensors, int64_t dim, at::Tensor & out) { at::TensorList tensors_; if (at::functionalization::impl::isFunctionalTensor(tensors)) { at::functionalization::impl::sync(tensors); tensors_ = at::functionalization::impl::from_functional_tensor(tensors); } else { tensors_ = tensors; } .... } ``` And here's the new one (the temporary it creates is a vector) ``` at::Tensor & cat_out_out(c10::DispatchKeySet dispatchKeySet, at::TensorList tensors, int64_t dim, at::Tensor & out) { ::std::vector<at::Tensor> tensors_; if (at::functionalization::impl::isFunctionalTensor(tensors)) { at::functionalization::impl::sync(tensors); tensors_ = at::functionalization::impl::from_functional_tensor(tensors); } else { tensors_ = tensors; } ... } ``` [ghstack-poisoned]

@ezyang

…g tensors" Addresses the comment here: #73441 (comment) This PR should also fix a bug with `at::cat.out` in functionalization, so I added a test for it. When the functionalization pass unwraps tensor arguments, we need to make sure the thing that we unwrap into has an owning type - we can't blindly base the type off whatever types come from the operator's schema. This is basically just a problem for `at::TensorList` (non-owning) - we need to use an owning `std::vector<Tensor>` instead. @ezyang I wasn't sure how general of a fix to make. `translate()` already knows a bit about "expensive conversions" from non-owning to owning types. I couldn't think of a way to re-use that code without making some larger changes (like teaching `CType`'s about whether they're owning types), so I added a simple fix inside the functionalization codegen. Feedback totally welcome. Here's part of the old `cat` kernel: ``` at::Tensor & cat_out_out(c10::DispatchKeySet dispatchKeySet, at::TensorList tensors, int64_t dim, at::Tensor & out) { at::TensorList tensors_; if (at::functionalization::impl::isFunctionalTensor(tensors)) { at::functionalization::impl::sync(tensors); tensors_ = at::functionalization::impl::from_functional_tensor(tensors); } else { tensors_ = tensors; } .... } ``` And here's the new one (the temporary it creates is a vector) ``` at::Tensor & cat_out_out(c10::DispatchKeySet dispatchKeySet, at::TensorList tensors, int64_t dim, at::Tensor & out) { ::std::vector<at::Tensor> tensors_; if (at::functionalization::impl::isFunctionalTensor(tensors)) { at::functionalization::impl::sync(tensors); tensors_ = at::functionalization::impl::from_functional_tensor(tensors); } else { tensors_ = tensors; } ... } ``` [ghstack-poisoned]

bdhirsh · 2022-04-25T21:32:46Z

@pytorchbot merge this please

pytorchmergebot · 2022-04-25T21:34:20Z

Merge failed due to Command git -C /home/runner/work/pytorch/pytorch push origin master returned non-zero exit code 1
To https://github.com/pytorch/pytorch
! [remote rejected] master -> master (cannot lock ref 'refs/heads/master': is at 40f3e85 but expected 5da76ac)
error: failed to push some refs to 'https://github.com/pytorch/pytorch'

Raised by https://github.com/pytorch/pytorch/actions/runs/2223007110

bdhirsh · 2022-04-25T21:47:17Z

@pytorchbot merge this please

pytorchmergebot · 2022-04-25T21:48:44Z

Merge failed due to Command git -C /home/runner/work/pytorch/pytorch cherry-pick -x 5a8f6814c4e9354929062574d6fd2e3f3ead5105 returned non-zero exit code 1
Auto-merging aten/src/ATen/native/native_functions.yaml
Auto-merging tools/autograd/gen_python_functions.py
CONFLICT (content): Merge conflict in tools/autograd/gen_python_functions.py
error: could not apply 5a8f681... functionalization: add a copy() native function
hint: After resolving the conflicts, mark them with
hint: "git add/rm ", then run
hint: "git cherry-pick --continue".
hint: You can instead skip this commit with "git cherry-pick --skip".
hint: To abort and get back to the state before "git cherry-pick",
hint: run "git cherry-pick --abort".

Raised by https://github.com/pytorch/pytorch/actions/runs/2223067256

@ezyang

…g tensors" Addresses the comment here: #73441 (comment) This PR should also fix a bug with `at::cat.out` in functionalization, so I added a test for it. When the functionalization pass unwraps tensor arguments, we need to make sure the thing that we unwrap into has an owning type - we can't blindly base the type off whatever types come from the operator's schema. This is basically just a problem for `at::TensorList` (non-owning) - we need to use an owning `std::vector<Tensor>` instead. @ezyang I wasn't sure how general of a fix to make. `translate()` already knows a bit about "expensive conversions" from non-owning to owning types. I couldn't think of a way to re-use that code without making some larger changes (like teaching `CType`'s about whether they're owning types), so I added a simple fix inside the functionalization codegen. Feedback totally welcome. Here's part of the old `cat` kernel: ``` at::Tensor & cat_out_out(c10::DispatchKeySet dispatchKeySet, at::TensorList tensors, int64_t dim, at::Tensor & out) { at::TensorList tensors_; if (at::functionalization::impl::isFunctionalTensor(tensors)) { at::functionalization::impl::sync(tensors); tensors_ = at::functionalization::impl::from_functional_tensor(tensors); } else { tensors_ = tensors; } .... } ``` And here's the new one (the temporary it creates is a vector) ``` at::Tensor & cat_out_out(c10::DispatchKeySet dispatchKeySet, at::TensorList tensors, int64_t dim, at::Tensor & out) { ::std::vector<at::Tensor> tensors_; if (at::functionalization::impl::isFunctionalTensor(tensors)) { at::functionalization::impl::sync(tensors); tensors_ = at::functionalization::impl::from_functional_tensor(tensors); } else { tensors_ = tensors; } ... } ``` [ghstack-poisoned]

bdhirsh · 2022-04-25T21:52:26Z

@pytorchbot merge this please

pytorchmergebot · 2022-04-25T21:53:44Z

Merge failed due to Refusing to merge as mandatory check Lint failed for rule superuser
Raised by https://github.com/pytorch/pytorch/actions/runs/2223089719

bdhirsh · 2022-04-25T21:58:45Z

@pytorchbot merge this please

github-actions · 2022-04-25T22:00:57Z

Hey @bdhirsh.
You've committed this PR, but it does not have both a 'release notes: ...' and 'topics: ...' label. Please add one of each to the PR. The 'release notes: ...' label should represent the part of PyTorch that this PR changes (fx, autograd, distributed, etc) and the 'topics: ...' label should represent the kind of PR it is (not user facing, new feature, bug fix, perf improvement, etc). The list of valid labels can be found here for the 'release notes: ...' and here for the 'topics: ...'.
For changes that are 'topic: not user facing' there is no need for a release notes label.

…76125) Summary: Pull Request resolved: #76125 Approved by: https://github.com/ezyang Test Plan: contbuild & OSS CI, see https://hud.pytorch.org/commit/pytorch/pytorch/640ce6bc9bbd25603c80651fd2be3f0045ea0d75 Reviewed By: osalpekar Differential Revision: D35938196 Pulled By: bdhirsh fbshipit-source-id: 1640a07e70f4036e42f12d657628e37848e26c35

functionalization bugfix: using owning type when unwrapping tensors

42dac3a

[ghstack-poisoned]

facebook-github-bot added the cla signed label Apr 20, 2022

bdhirsh commented Apr 20, 2022

View reviewed changes

bdhirsh requested review from ezyang and ysiraichi April 20, 2022 19:05

bdhirsh mentioned this pull request Apr 20, 2022

make functionalization work better with subclasses #73441

Closed

ezyang approved these changes Apr 21, 2022

View reviewed changes

This was referenced Apr 21, 2022

functionalization: fix nested grad + functionalize transforms #76189

Closed

remove _is_foreach_op codegen special cases, clean up mutable return type checks #76190

Closed

bdhirsh added 4 commits April 21, 2022 16:36

This was referenced Apr 25, 2022

fix nested grad(functionalize(f)) transforms #76318

Closed

fix torch.tensor for functionalization #76319

Closed

generate out= and functional variants of NativeFunctions, get functionalization to work for all mutable ops #76320

Closed

pytorchmergebot closed this in 640ce6b Apr 25, 2022

facebook-github-bot deleted the gh/bdhirsh/213/head branch April 29, 2022 14:17

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

functionalization bugfix: using owning type when unwrapping tensors#76125

functionalization bugfix: using owning type when unwrapping tensors#76125
bdhirsh wants to merge 8 commits intogh/bdhirsh/213/basefrom
gh/bdhirsh/213/head

bdhirsh commented Apr 20, 2022 •

edited

Loading

Uh oh!

facebook-github-bot commented Apr 20, 2022 •

edited

Loading

Uh oh!

bdhirsh Apr 20, 2022

Uh oh!

ezyang left a comment

Uh oh!

bdhirsh commented Apr 25, 2022

Uh oh!

pytorchmergebot commented Apr 25, 2022

Uh oh!

bdhirsh commented Apr 25, 2022

Uh oh!

pytorchmergebot commented Apr 25, 2022

Uh oh!

bdhirsh commented Apr 25, 2022

Uh oh!

pytorchmergebot commented Apr 25, 2022

Uh oh!

bdhirsh commented Apr 25, 2022

Uh oh!

github-actions bot commented Apr 25, 2022

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Conversation

bdhirsh commented Apr 20, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

facebook-github-bot commented Apr 20, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful links

💊 CI failures summary and remediations

Uh oh!

bdhirsh Apr 20, 2022

Choose a reason for hiding this comment

Uh oh!

ezyang left a comment

Choose a reason for hiding this comment

Uh oh!

bdhirsh commented Apr 25, 2022

Uh oh!

pytorchmergebot commented Apr 25, 2022

Uh oh!

bdhirsh commented Apr 25, 2022

Uh oh!

pytorchmergebot commented Apr 25, 2022

Uh oh!

bdhirsh commented Apr 25, 2022

Uh oh!

pytorchmergebot commented Apr 25, 2022

Uh oh!

bdhirsh commented Apr 25, 2022

Uh oh!

github-actions bot commented Apr 25, 2022

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

bdhirsh commented Apr 20, 2022 •

edited

Loading

facebook-github-bot commented Apr 20, 2022 •

edited

Loading