[inductor] Add lowering for as_strided_scatter

[peterbell10]

Stack from ghstack (oldest at bottom):

• [primTorch] Add prim and ref for as_strided_scatter #88426
• -> [inductor] Add lowering for as_strided_scatter #88379

Ref #93757

The use of as_strided does require in-memory manipulations, however this
lowering allows those memory ops to be fused with any preceding calculations.
e.g.

def f(a, b):
    return torch.as_strided_scatter(
        a * 8 + 10,
        b * 2 - 4,
        size=(a.numel() // 2,),
        stride=(2,))

Before this compiles to two kernels and a call to aten.as_strided_scatter and
with this PR it compiles to just two kernels and no additional operator calls.

In theory I think this could be a decomposition, but in practice I saw the
output_view.copy_(src) being optimized out in some cases when this was
implemented as a decomposition.

cc @mlazos @soumith @voznesenskym @yanboliang @penguinwu @anijain2305 @EikanWang @jgong5 @Guobing-Chen @chunyuan-w @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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

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

❌ 1 Failures, 1 Pending

As of commit 02a819e:

The following jobs have failed:

• Check labels

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

[lezcano]

Note that it'd still be valuable to implement this as a decomposition for other backends to use.

[jansel]

Actually, can you add a test for this?

[peterbell10]

@jansel PTAL, I've added a test.

[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

[ngimel]

output_view.copy_(src)

being optimized away is worrying, @peterbell10 do you have a repro? cc @bdhirsh

[peterbell10]

On the next PR in the stack (edit: specifically at commit hash deca398), if I enable the decomposition in torch/_inductor/decomposition.py then this reproducer fails for me:

import torch
from torch.fx.experimental.proxy_tensor import make_fx
from torch._inductor.compile_fx import compile_fx

x = torch.rand((1,), device="cuda", requires_grad=True)
y = torch.rand((1,), device="cuda", requires_grad=True)
def f(a, b):
    return torch.as_strided_scatter(
        a,
        b,
        size=(1,),
        stride=(1,)),

torch._inductor.config.debug = True
args = [x, y]

decomposed = make_fx(f, tracing_mode="fake")(*args)
compiled_decomposed = compile_fx(decomposed, args)
expect = f(*args)
actual = compiled_decomposed(*args)
print(f'Expected {expect}\nActual: {actual}')

The requires_grad=True on the input is important.

[bdhirsh]

Thanks for the repro @peterbell10

I'm going to add asserts in aot autograd soon to check that there are no mutable ops in the graph right before we run DCE, which is what causes problems.

For the actual fix, I think what we should probably do is use functionalization to "functionalize" every decomposition right before we run it, during proxy tensor tracing. That way we don't need to worry about / establish invariants on whether decomps are allowed to call mutations. I'm gonna try to POC this soon.

There's probably some question as to whether or not proxy tensor tracing should functionalize decomps by default or not. Maybe we need another feature flag, in make_fx? (lol)