Read out real strides from compilation result, rather than real args by ezyang · Pull Request #105010 · pytorch/pytorch

ezyang · 2023-07-11T21:36:35Z

Stack from ghstack (oldest at bottom):

This prefigures a refactor that will move the backward compilation
to entirely ahead of time, so I need to extract these strides some
other way. Straight from the compiler's mouth will do it.

I can't easily get the information via the return result of fw_compiler without changing the calling convention, so instead I smuggle it via TracingContext. TracingContext may be None when we are compiling patterns for the joint graph pattern matcher.

Signed-off-by: Edward Z. Yang ezyang@meta.com

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @ngimel @yf225 @chenyang78 @kadeng @muchulee8

This prefigures a refactor that will move the backward compilation to entirely ahead of time, so I need to extract these strides some other way. Straight from the compiler's mouth will do it. Signed-off-by: Edward Z. Yang <ezyang@meta.com> [ghstack-poisoned]

pytorch-bot · 2023-07-11T21:36:38Z

🔗 Helpful Links

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

📄 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 e2a8ecd:
💚 Looks good so far! There are no failures yet. 💚

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

shunting314 · 2023-07-11T23:21:51Z

torch/_guards.py

+        # compiler to aot_autograd
+        # Per output, what the compiler specified stride of the output is,
+        # or None if no stride is known
+        self.output_strides: Optional[List[Optional[List[int]]]] = None


Right now we do layout optimization only if dynamic shape is disabled. But we are working on resolving that (blocked on the split reduction support). So, should we set the type here to consider SymInt?

yeah this can have symints, will amend

shunting314 · 2023-07-11T23:28:23Z

torch/_inductor/compile_fx.py

+                # Return the output strides to the caller via TracingContext
+                assert len(context.output_strides) == 0
+                for out in graph.graph_outputs:
+                    if hasattr(out, "layout"):


Do we see cases that out does not have a layout attribute? Does this happen for some SymInt returned?

yes, can have symint return. Maybe there is something better than hasattr to do here, but a lot of inductor code is written this way. Hard to be more precise without typing 👹

OK, turns out I can't, because strides in inductor are sympy.Symbol, not SymInt. So I am just going ahead and storing the hints here only. When you fix this to do permutations instead, probably can make this a little less fragile.

We can not do permutations instead because of some non 'dense' activations? This happens for some real models.

There are a few ways we can do this. If we only ever change layout on dense tensors, we can make output_strides be None unless we changed the stride. Then the permutation is always defined.

If we only ever change layout on dense tensors, ...

I think this can not be guaranteed in inductor right now. A non dense tensor's layout may get changes because of the layout change of its upstream tensors. And inductor does not force eager stride for that non dense tensor because of code here: https://github.com/pytorch/pytorch/blob/main/torch/_inductor/graph.py#L688 . I guess the algorithm will be tricky to restride a non-dense tensor.

OK will need to think this through carefully. I'll start a doc.

Upon further reflection, all of this is moot if we have compiled backwards, so let's not touch it unless it's causing someone problems.

shunting314 · 2023-07-11T23:32:51Z

torch/_guards.py

+        if tc is None:
+            yield None
+            return
+        old_output_strides = tc.output_strides


Just want to raise a concern about nested compiling where we may need a stack for output_strides. But I can not come up with a realistic example. So probably the current implementation is good enough

This implicitly is a stack via old_output_strides!

ah, right, that's the whole point of context manager after all... haha

shunting314 · 2023-07-11T23:36:01Z

torch/_functorch/aot_autograd.py

                            continue

-                        # Comparing ph_arg.stride() with real_arg.stride() directly may
+                        if forward_saved_for_backwards_strides is None:


Move this check out of the for loop may be slightly faster.

I was trying to save myself an indent 😂 This is only during compilation so I don't think it matters much

shunting314 · 2023-07-11T23:37:33Z

torch/_functorch/aot_autograd.py

+                        if real_stride is None:
+                            continue
+
+                        assert _get_hints(real_stride) == all_args[i].stride(), f"{real_stride} {all_args[i].stride()}"


I assume you will remove this soon since the main point of the change is we won't have access to all_args?

Yup, this is just for CI here.

shunting314 · 2023-07-11T23:44:47Z

torch/_functorch/aot_autograd.py

+                        forward_saved_for_backwards_strides = fwd_output_strides[
+                            CompiledFunction.metadata.tensors_saved_for_backwards_slice
+                        ]


This looks nice!

… real args" This prefigures a refactor that will move the backward compilation to entirely ahead of time, so I need to extract these strides some other way. Straight from the compiler's mouth will do it. I can't easily get the information via the return result of `fw_compiler` without changing the calling convention, so instead I smuggle it via TracingContext. TracingContext may be None when we are compiling patterns for the joint graph pattern matcher. Signed-off-by: Edward Z. Yang <ezyangmeta.com> cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx peterbell10 ipiszy ngimel yf225 chenyang78 kadeng muchulee8 [ghstack-poisoned]

This prefigures a refactor that will move the backward compilation to entirely ahead of time, so I need to extract these strides some other way. Straight from the compiler's mouth will do it. Signed-off-by: Edward Z. Yang <ezyangmeta.com> ghstack-source-id: f9ceb24 Pull Request resolved: #105010

ezyang · 2023-07-12T11:30:56Z

@pytorchbot merge

pytorchmergebot · 2023-07-12T11:33:01Z

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

…05251) Currently all information about the dependencies of ghstack PRs (e.g. #105010) is stripped away: https://github.com/pytorch/pytorch/blob/c984885809194e0a807b3f5543450fae4dfa841a/.github/scripts/trymerge.py#L1077-L1078 This PR adds this information back in a more compact form. All dependencies (PR numbers) of each PR in ghstack are recorded. The resulting commit message will look like this (the last line is new): > Mock title (#123) > > Mock body text > Pull Request resolved: #123 > Approved by: https://github.com/Approver1, https://github.com/Approver2 > ghstack dependencies: #1, #2 --- ### Testing Unit tests. --- ### Note Re: `# type: ignore[assignment]` in unit tests. I did my due diligence to find alternatives. Unfortunately mypy [doesn't](python/mypy#6713) support this [way of patching methods](https://docs.python.org/3/library/unittest.mock-examples.html#mock-patching-methods), and the alternatives are either extremely verbose or don't work for this case. I decided it's not worth the effort (since the problem is limited only to the unit test). Pull Request resolved: #105251 Approved by: https://github.com/huydhn

ezyang mentioned this pull request Jul 11, 2023

Move more stuff into ViewAndMutationMeta #105009

Closed

pytorch-bot bot added the release notes: AO frontend label Jul 11, 2023

ezyang mentioned this pull request Jul 11, 2023

Immediately compile backwards graph in AOTAutograd if dynamic shapes #104971

Closed

github-actions bot added module: inductor ciflow/inductor labels Jul 11, 2023

github-actions bot requested review from SherlockNoMad, albanD, antoniojkim, bdhirsh, jbschlosser, miladm, voznesenskym and wconstab July 11, 2023 21:36

albanD removed their request for review July 11, 2023 21:55

ezyang requested a review from shunting314 July 11, 2023 23:13

ezyang added the ciflow/trunk Trigger trunk jobs on your pull request label Jul 11, 2023

shunting314 reviewed Jul 11, 2023

View reviewed changes

shunting314 approved these changes Jul 11, 2023

View reviewed changes

pytorchmergebot added the merging label Jul 12, 2023

pytorchmergebot added Merged and removed merging labels Jul 12, 2023

pytorchmergebot closed this in 979f826 Jul 12, 2023

izaitsevfb mentioned this pull request Jul 14, 2023

[GHF][mergebot] record ghstack dependencies in the commit message #105251

Closed

facebook-github-bot deleted the gh/ezyang/2222/head branch July 15, 2023 14:16

Conversation

ezyang commented Jul 11, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Jul 11, 2023 • 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/105010

✅ No Failures

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

shunting314 Jul 11, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ezyang commented Jul 12, 2023

Uh oh!

pytorchmergebot commented Jul 12, 2023

Merge started

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

ezyang commented Jul 11, 2023 •

edited

Loading

pytorch-bot bot commented Jul 11, 2023 •

edited

Loading

shunting314 Jul 11, 2023 •

edited

Loading