Avoid generating as_strided for alaising views in auto_functionalize_v2

[laithsakka]

Stack from ghstack (oldest at bottom):

• -> Avoid generating as_strided for alaising views in auto_functionalize_v2 #137149

during auto_functionalize_v2 if we encounter a view such that size() stride() and storage_offset() matches the base
we create a view that is regenerated by calling aten.alias instead of as_strided for better performance.

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

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 ddba772 with merge base 839d356 ():
💚 Looks good so far! There are no failures yet. 💚

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

[zou3519]

during auto_functionalize_v2 if we encounter a view such that size() stride() and storage_offset() matches the base
we just pass the base. (we can potentially call alias instead to be persisitance with weried case that check id inside custom op) not sure if we shall do that.

Can the Tensor ever be used later on? If so then we need to call alias to preserve the semantics (that the tensor is a view instead of the original tensor)

[laithsakka]

during auto_functionalize_v2 if we encounter a view such that size() stride() and storage_offset() matches the base
we just pass the base. (we can potentially call alias instead to be persisitance with weried case that check id inside custom op) not sure if we shall do that.

Can the Tensor ever be used later on? If so then we need to call alias to preserve the semantics (that the tensor is a view instead of the original tensor)

inside the custom op?

[zou3519]

Let's beef up the testing with dynamic shapes. Also, we shouldn't delete the as_strided calls, they should be turned into alias calls.

[laithsakka]

so we do add guards when we fail the check
for example:

@torch.compile(fullgraph=True, dynamic=True, backend="inductor")
def f(a):
   b = a[0]
   foo(b)

+- LAMBDA_GUARD: Ne(L['a'].size()[0], L['a'].size()[2])  # (_higher_order_ops/auto_functionalize.py:81 in write_single_view)

we do not add new guards for this

@torch.compile(fullgraph=True, dynamic=True, backend="inductor")
def f(a):
   b = torch.ops.aten.alias (a)
   foo(b)

because the two side of the symbol we are comparing are the same symbols.

however what is the likely hood of actually failing the guards in next iterations, is your concern about the guard checking
over head that we can avoid if we do this statically (capturing the alias call). or is your concern with potential recompilation.

[laithsakka]

I see the point that "this guards is not needed in the case above, because we always know that it will is true, and we can write the opt in a way that does not generate the guard".

but someone can counter that with " well if its always the same then we dont have to worry about recompilation so unless we are worried about guard checking time we shall be ok"

I do not have strong opinion about landing this just sharing my thoughts, lmk what do you think

[zou3519]

The as_strided -> alias() change sounds good then, because it does not add new guards. Can you add additional tests to show that it doesn't invoke a recompile when the shape changes? Then we can land this PR.

I'm just concerned about unnecessary recompiles.

For the slice change: I'm confused, why is it generating a guard that looks like Ne(L['a'].size()[0], L['a'].size()[2]) ?

[laithsakka]

as_strided -> alias()

oh Ne(L['a'].size()[0], L['a'].size()[2]) is added by the alias change not by the slice change. I did not check anything on the other slice change yet.
basically we fail the check

     base.size() == tensor.size() 

and this cause the above failure.

[laithsakka]

added recompilation tests, not sure if that covers your intention of :

we should also have a test where we create a tensor with unbacked symints, and then we take an alias of that, and then we mutate it.

also update the diff to generate alias, (added a test but it fail seems that soemwhere down the road inductor remove the alias ) see issue #137434

[zou3519]

To check... does inductor remove the copy_ during lowering?

[laithsakka]

output code

def call(args):
    arg0_1, arg1_1, arg2_1 = args
    args.clear()
    s0 = arg0_1
    s1 = arg1_1
    assert_size_stride(arg2_1, (s0, s1), (s1, 1))
    # Topologically Sorted Source Nodes: [], Original ATen: []
    torch.ops.mylib.foo.default(arg2_1, arg2_1)
    return (arg2_1, arg2_1, )

[zou3519]

can you make this something like...

def f(x):
    a = torch.ops.aten.alias.default(x)
    b = x.clone()
    c = b.nonzero().float()
    d = c.alias()
    torch.ops.mylib.foo(a, d)
    return a, d

d is a Tensor with unbacked symints in the shape

[laithsakka]

i will add a test for that i here and in the recompile test

[laithsakka]

added they pass

[zou3519]

hmmm

[zou3519]

can you return a? and then assert that the input and output of func have different id identities.

Inductor doesn't match the tensor identity of intermediates of the function, but it should better match the identity of inputs/outputs of the compiled function.

[laithsakka]

i will add another test that does that

[laithsakka]

added that , that one pass.

[zou3519]

code LGTM, but please see my suggestions for the test cases

[laithsakka]

@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