[AOTDispatch] Return mutated inputs directly when keeping mutations

[peterbell10]

Stack from ghstack (oldest at bottom):

• -> #120514

Fixes #120242

The example from the issue now results in the graph

def forward(self, arg0_1, arg1_1):
    sin = torch.ops.aten.sin.default(arg0_1);  arg0_1 = None
    copy_ = torch.ops.aten.copy_.default(arg1_1, sin);  arg1_1 = sin = None
    return (arg0_1,)

and the corresponding inductor kernel eliminates the intermediate buffer completely

def call(args):
    arg0_1, arg1_1 = args
    args.clear()
    assert_size_stride(arg0_1, (5, ), (1, ))
    assert_size_stride(arg1_1, (5, ), (1, ))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        # Source Nodes: [sin], Original ATen: [aten.sin]
        stream0 = get_raw_stream(0)
        triton_poi_fused_sin_0.run(arg0_1, arg1_1, 5, grid=grid(5), stream=stream0)
        del arg0_1
    return (arg1_1, )

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

[pytorch-bot[bot]]

:link: Helpful Links

:test_tube: See artifacts and rendered test results at hud.pytorch.org/pr/120514

• :page_facing_up: Preview Python docs built from this PR
• :page_facing_up: Preview C++ docs built from this PR
• :question: 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.

:white_check_mark: No Failures

As of commit 39c69f33323b2b1f8cb008da5489d2bcc5aff527 with merge base 953c6c37cbfcc894963120666dbf066977931885 (): :green_heart: Looks good so far! There are no failures yet. :green_heart:

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

[lezcano]

ping @bdhirsh

[ezyang]

@bdhirsh is on vacation for two weeks.

[lezcano]

Let's just wait for @bdhirsh to be back then, as this issue is not blocking anything.

[ezyang]

I'm 👍 the AOTAutograd changes. But the Inductor passes also need reviewing. I can bug @Chillee / @oulgen to look at it in person, or in a pinch I can review them too.

[peterbell10]

@pytorchbot merge -r

[pytorchmergebot]

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

[pytorchmergebot]

Rebase failed due to Command git -C /home/runner/work/pytorch/pytorch rebase refs/remotes/origin/viable/strict gh/peterbell10/697/orig returned non-zero exit code 1

Rebasing (1/1)
Auto-merging test/functorch/test_aotdispatch.py
Auto-merging torch/_functorch/_aot_autograd/traced_function_transforms.py
Auto-merging torch/_inductor/fx_passes/post_grad.py
Auto-merging torch/_inductor/fx_passes/reinplace.py
CONFLICT (content): Merge conflict in torch/_inductor/fx_passes/reinplace.py
error: could not apply 8a18644291f... [AOTDispatch] Return mutated inputs directly when keeping mutations
hint: Resolve all conflicts manually, mark them as resolved with
hint: "git add/rm <conflicted_files>", then run "git rebase --continue".
hint: You can instead skip this commit: run "git rebase --skip".
hint: To abort and get back to the state before "git rebase", run "git rebase --abort".
Could not apply 8a18644291f... [AOTDispatch] Return mutated inputs directly when keeping mutations

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

[peterbell10]

@pytorchbot merge

[pytorchmergebot]

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