AOTAutograd perf: avoid as_strided() calls when we have intermediate bases

[bdhirsh]

This is a more targeted version of an existing issue around as_strided calls in AOTAutograd, #109237. Came from an internal issue

Simple repro:

import torch

@torch.compile
def f(x):
    out = x.mul(2)
    return out.view(out.shape), out.view(out.shape)

inps = (torch.randn(5), torch.tensor(0))
x = torch.randn(4, requires_grad=True)
out1, out2 = f(x)
print(out1.grad_fn)

prints:

<AsStridedBackward0 object at 0x7f4c7d4d2260>

We end up calling as_strided in the compiled forward, so an AsStridedBackward node shows up in the backward, which in general is not implemented to be particularly fast.

Why does this happen?

(1) AOTAutograd has logic for "intermediate bases". If we have two outputs of our graph that are aliases of each other (and of the same graph intermediate), today, AOTAutograd will just have the shared intermediate be an output to the graph. AOTAutograd will then replay the views off of the intermediate, so that autograd properly realizes that the outputs alias.

(2) AOTAutograd has a function to try to do the view replay, but it hits a slow path in that function that causes it to go to as_strided. We should figure out why and fix this: https://github.com/pytorch/pytorch/blob/main/torch/_functorch/aot_autograd.py#L807

cc @ezyang @msaroufim @wconstab @anijain2305 @zou3519

[bdhirsh]

In fact, a different (but simpler) example is this - no intermediate bases, just an output that aliases the input

import torch

@torch.compile(backend="inductor")
def f(x):
    return x.view(x.shape)

x = torch.randn(4, requires_grad=True)
out1 = f(x)
print(out1.grad_fn)

<AsStridedBackward0 object at 0x7f87b84d6c50>

[bdhirsh]

We need to properly replay these views that happened in the graph, without relying on autograd's view replay.

The most obvious way to do this is probably: functionalization already tracks these views somewhere, so we can ask functionalization to replay them (this should also be pretty fast, since it's all stored in lambdas in C++).

One small downside to doing it this way, though, is that I don't think we can use this approach when input are subclasses (and the output of the compiled blob is a view of that subclass). Since functionalization runs below the subclass, and the subclass might insert other logic around the view.

We can always deal with that problem later though.

[bdhirsh]

I'm testing out the approach above.

This approach doesn't work with dynamic shapes: functionalization does remember all the views that were used to generate the output, but it also stashes the symbolic values that were passed to each view op. In order to actually re-use functionalization's stashed view information, we'd need to resolve each symbol at runtime. This seems doable, but not trivial.

[gchanan]

Do you have a benchmark or similar you are tracking to ensure doing functionalization is worth it here?

[ezyang]

doc https://docs.google.com/document/d/1DlfFq8TKbuAn2zyJxLfoW-X1qkkm5PLdHFtySo03QAk/edit

[bdhirsh]

Do you have a benchmark or similar you are tracking to ensure doing functionalization is worth it here?

Talked offline - agreed that if we switch from as_strided to view chains as the default in all cases, we should do some amount of benchmarking to make sure we aren't regressing any cases.

[bdhirsh]

#111411 should fix the most egregious instances of as_strided. It turns out that by-far the worst offenders are multi-output views. Take these two examples:

(alias-of-intermediate case)

def f(x):
    intermediate = x * x
    out1, out2, ... out50 = intermediate.unbind(0)
    return out1, out2, ... out50

(alias-of-input case)

def f(x):
    intermediate = x * x
    out1, out2, ... out50 = intermediate.unbind(0)
    return out1, out2, ... out50

It turns out that multi-output views are exactly the case that autograd's view-replay logic cannot handle, so we always fall back to as_strided in both of these cases.

Worst of all, the original code had a single call to UnbindBackward in the autograd graph, but the AOTAutograd code here will have 50 AsStridedBackward nodes in the backward graph (one for each output, since we call as_strided individually on each output.

The "partition the forward graph across view outputs" idea from this doc https://docs.google.com/document/d/1DlfFq8TKbuAn2zyJxLfoW-X1qkkm5PLdHFtySo03QAk/edit is a good long-term fix for this problem. But it's very far out - it requires some major surgery to AOTAutograd, and isn't viable until pre-dispatch functionalization is both completed, and hardened enough to run 100% of the time in AOTAutograd (cc @tugsbayasgalan 😃).

Instead, the approach that I'm attempting to take in the linked PR is to effectively hide all multi-output-view aliasing from autograd. We should think carefully about whether or not this is safe in all cases. But it relies on the high-level idea that the autograd engine does not allow you to mutate the outputs of multi-output views:

x = torch.ones(2, 2, requires_grad=True).clone()
y, z = x.unbind(0)

y.mul_(2)  # This raises an autograd error, cannot mutate the output of a multi-output-view!
x.mul_(2) # This works, BUT swaps out y.grad_fn and z.grad_fn with error nodes!