Fix FallbackKernel behavior on mutable ops #118649

zou3519 · 2024-01-30T17:30:23Z

Stack from ghstack (oldest at bottom):

FallbackKernel wasn't handing mutable ops correctly: it would not report
them in get_mutation_names or get_alias_names. This would lead to silent
incorrectness -- Inductor would incorrectly reorder the mutable op with other
mutable ops.

This PR fixes that:

we only support mutable operations that are "auto_functionalizable".
That is, they mutate inputs and do not return aliases of any inputs.
Following the Triton kernel work, any mutated inputs must be specified
in get_alias_names and processed via mark_node_as_mutating
We also do some minor cleanup by killing dead code (FallbackKernel no
longer processes OpOverloadPacket) and adding some handling around
HOPs.

Test Plan:

new tests

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

FallbackKernel wasn't handing mutable ops correctly: it would not report them in get_mutation_names or get_alias_names. This would lead to silent incorrectness -- Inductor would incorrectly reorder the mutable op with other mutable ops. This PR fixes that: - we only support mutable operations that are "auto_functionalizable". That is, they mutate inputs and do not return aliases of any inputs. - Following the Triton kernel work, any mutated inputs must be specified in get_alias_names and processed via mark_node_as_mutating - We also do some minor cleanup by killing dead code (FallbackKernel no longer processes OpOverloadPacket) and adding some handling around HOPs. Test Plan: - new tests [ghstack-poisoned]

pytorch-bot · 2024-01-30T17:30:27Z

🔗 Helpful Links

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

📄 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 86e5690 with merge base 064610d ():
💚 Looks good so far! There are no failures yet. 💚

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

FallbackKernel wasn't handing mutable ops correctly: it would not report them in get_mutation_names or get_alias_names. This would lead to silent incorrectness -- Inductor would incorrectly reorder the mutable op with other mutable ops. This PR fixes that: - we only support mutable operations that are "auto_functionalizable". That is, they mutate inputs and do not return aliases of any inputs. - Following the Triton kernel work, any mutated inputs must be specified in get_alias_names and processed via mark_node_as_mutating - We also do some minor cleanup by killing dead code (FallbackKernel no longer processes OpOverloadPacket) and adding some handling around HOPs. Test Plan: - new tests ghstack-source-id: e85d925 Pull Request resolved: #118649

FallbackKernel wasn't handing mutable ops correctly: it would not report them in get_mutation_names or get_alias_names. This would lead to silent incorrectness -- Inductor would incorrectly reorder the mutable op with other mutable ops. This PR fixes that: - we only support mutable operations that are "auto_functionalizable". That is, they mutate inputs and do not return aliases of any inputs. - Following the Triton kernel work, any mutated inputs must be specified in get_alias_names and processed via mark_node_as_mutating - We also do some minor cleanup by killing dead code (FallbackKernel no longer processes OpOverloadPacket) and adding some handling around HOPs. Test Plan: - new tests cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx peterbell10 ipiszy yf225 chenyang78 kadeng muchulee8 aakhundov ColinPeppler [ghstack-poisoned]

eellison · 2024-01-31T21:00:07Z

torch/_inductor/ir.py

+        if isinstance(self.op_overload, torch._ops.HigherOrderOperator):
+            # We assume here that HOPs with FallbackKernel are functional.
+            # This may not always be true! HOPs must individually opt-in to
+            # FallbackKernel, so please check this if you opt-in.
+            return


Can we check this assumption somehow?

No, there's no notion of a schema for HOPs. Though we generally tell people that they must be functional.

torch/_library/utils.py

eellison · 2024-01-31T21:34:28Z

torch/_inductor/ir.py

+            return
+
+        if schema.is_mutable and not can_auto_functionalize(kernel):
+            raise NotImplementedError(


Why cant we handle this case ?

For custom ops, the schema is not enough information to tell us about the semantics of the operator. For example, if the operator has a schema that looks like "(Tensor(a!) x) -> Tensor(a)", it's unclear if the operator is (1) returning x as-is or (2) returning a view on x.

For aten operators, we have the convention that it is always (1). It's unclear to me how to generate code when we don't know which of these two situations we're in.

Pragmatically, I've never seen a custom op have a schema that looks like the above, so, this is not a case we need to worry about.

eellison · 2024-01-31T21:41:22Z

torch/_inductor/ir.py

+        schema_args = schema.arguments
+        args, kwargs = self.unflatten_args(self.inputs, self.constant_args)
+        input_alias_sets = set()
+        for arg, info in zip(args, schema_args):


I'm not sure that these two match up. schema_args will not unroll Tensor lists for example. could we add some assertions of some kind or another

Yeah, I'll work through this more

I think an assertion for isinstance(Tensor) is sufficient here, will put it in (and some comments about why)

What about func(Tensor[] inp, Tensor (a!) inp2) ? Wont we still unflatten the inp list and have the mutation arg not be aligned ?

We've got a test for this. But to answer your question, self.unflatten_args puts the flat list back into the structure that is aligned with the schema. Though I realized I forgot to handle kwargs, will go do that.

@eellison I put in a bunch of safety checks (and we handle kwargs now), please let me know if that's what you were looking for.

torch/_inductor/ir.py

FallbackKernel wasn't handing mutable ops correctly: it would not report them in get_mutation_names or get_alias_names. This would lead to silent incorrectness -- Inductor would incorrectly reorder the mutable op with other mutable ops. This PR fixes that: - we only support mutable operations that are "auto_functionalizable". That is, they mutate inputs and do not return aliases of any inputs. - Following the Triton kernel work, any mutated inputs must be specified in get_alias_names and processed via mark_node_as_mutating - We also do some minor cleanup by killing dead code (FallbackKernel no longer processes OpOverloadPacket) and adding some handling around HOPs. Test Plan: - new tests cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx peterbell10 ipiszy yf225 chenyang78 kadeng muchulee8 aakhundov ColinPeppler [ghstack-poisoned]

oulgen · 2024-02-08T19:34:24Z

torch/_inductor/ir.py

            return devices[0]
        return None

    def has_side_effects(self):


Unrelated to this PR: I want us to kill has_side_effects. We should detect this, not mark it.

cc: @eellison @Chillee

oulgen · 2024-02-08T19:35:07Z

torch/_inductor/ir.py

+        return self.alias_names
+
+    def get_mutation_names(self):
+        return self.mutation_names


do we need to assert that this 0 or 1 elements?

Yes, will do

FallbackKernel wasn't handing mutable ops correctly: it would not report them in get_mutation_names or get_alias_names. This would lead to silent incorrectness -- Inductor would incorrectly reorder the mutable op with other mutable ops. This PR fixes that: - we only support mutable operations that are "auto_functionalizable". That is, they mutate inputs and do not return aliases of any inputs. - Following the Triton kernel work, any mutated inputs must be specified in get_alias_names and processed via mark_node_as_mutating - We also do some minor cleanup by killing dead code (FallbackKernel no longer processes OpOverloadPacket) and adding some handling around HOPs. Test Plan: - new tests cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx peterbell10 ipiszy yf225 chenyang78 kadeng muchulee8 aakhundov ColinPeppler [ghstack-poisoned]

zou3519 · 2024-02-09T18:59:17Z

@pytorchbot merge

pytorchmergebot · 2024-02-09T19:01:18Z

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

FallbackKernel wasn't handing mutable ops correctly: it would not report them in get_mutation_names or get_alias_names. This would lead to silent incorrectness -- Inductor would incorrectly reorder the mutable op with other mutable ops. This PR fixes that: - we only support mutable operations that are "auto_functionalizable". That is, they mutate inputs and do not return aliases of any inputs. - Following the Triton kernel work, any mutated inputs must be specified in get_alias_names and processed via mark_node_as_mutating - We also do some minor cleanup by killing dead code (FallbackKernel no longer processes OpOverloadPacket) and adding some handling around HOPs. Test Plan: - new tests Pull Request resolved: #118649 Approved by: https://github.com/eellison, https://github.com/oulgen

github-actions bot added module: inductor ciflow/inductor labels Jan 30, 2024

zou3519 added the keep-going Don't stop on first failure, keep running tests until the end label Jan 30, 2024

zou3519 mentioned this pull request Jan 30, 2024

Decompose torch.ops.higher_order.auto_functionalized in Inductor #118673

Closed

zou3519 added 3 commits January 30, 2024 17:59

zou3519 requested review from bdhirsh, eellison and oulgen January 31, 2024 18:13

eellison reviewed Jan 31, 2024

View reviewed changes

bdhirsh reviewed Jan 31, 2024

View reviewed changes

torch/_inductor/ir.py Show resolved Hide resolved

zou3519 requested a review from eellison February 5, 2024 21:18

zou3519 added 2 commits February 6, 2024 06:56

zou3519 requested a review from Chillee February 6, 2024 22:20

zou3519 added 2 commits February 6, 2024 18:58

eellison approved these changes Feb 8, 2024

View reviewed changes

zou3519 added ciflow/trunk Trigger trunk jobs on your pull request release notes: inductor labels Feb 8, 2024

oulgen reviewed Feb 8, 2024

View reviewed changes

oulgen approved these changes Feb 8, 2024

View reviewed changes

pytorchmergebot added the merging label Feb 9, 2024

pytorchmergebot closed this in 01e248d Feb 9, 2024

pytorchmergebot added Merged and removed merging labels Feb 9, 2024

github-actions bot deleted the gh/zou3519/900/head branch March 11, 2024 01:52

Fix FallbackKernel behavior on mutable ops #118649

Fix FallbackKernel behavior on mutable ops #118649

Uh oh!

Conversation

zou3519 commented Jan 30, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Jan 30, 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/118649

✅ No Failures

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

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!

zou3519 Feb 5, 2024 • 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!

zou3519 Feb 7, 2024 • 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!

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!

zou3519 commented Feb 9, 2024

Uh oh!

pytorchmergebot commented Feb 9, 2024

Merge started

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

zou3519 commented Jan 30, 2024 •

edited

Loading

pytorch-bot bot commented Jan 30, 2024 •

edited

Loading

zou3519 Feb 5, 2024 •

edited

Loading

zou3519 Feb 7, 2024 •

edited

Loading