dynamo.optimizations.training.aot_autograd does not functionalize

[silvasean]

🐛 Describe the bug

I tried adjusting my code from using raw make_fx to aot_autograd as suggested by @ezyang in this comment. It seems to be working well for converting the graph to dispatcher ops, and gives us the backward graph too as Ed said. However, Ed mentioned that it would also give us functionalization, but I don't see that happening. I tried poking around inside aot_autograd to force it down the normalize_ir code path and didn't get the functionalized graph either.

See repro script and output:

import torch
import torch._dynamo as dynamo
from torch._dynamo.optimizations.training import aot_autograd
import functorch

@functorch.compile.make_boxed_compiler
def my_backend(gm, example_inputs):
    gm.graph.print_tabular()
    return gm

@dynamo.optimize(aot_autograd(fw_compiler=my_backend))
def f(x):
    x[0].relu_()
    return x

m1 = -torch.ones(2, 5)
result = f(m1)
print("m1", m1)
print("result", result)
print("m1 is result:", m1 is result)

Output:

opcode         name    target              args            kwargs
-------------  ------  ------------------  --------------  --------
placeholder    arg0_1  arg0_1              ()              {}
call_function  select  aten.select.int     (arg0_1, 0, 0)  {}
call_function  relu_   aten.relu_.default  (select,)       {}
output         output  output              ((arg0_1,),)    {}
m1 tensor([[ 0.,  0.,  0.,  0.,  0.],
        [-1., -1., -1., -1., -1.]])
result tensor([[ 0.,  0.,  0.,  0.,  0.],
        [-1., -1., -1., -1., -1.]])
m1 is result: True

Compare this to the result with raw use of make_fx/functionalize:

import torch
import torch._dynamo as dynamo
import functorch
from functorch.experimental import functionalize

def my_backend(gm, example_inputs):
    # Functionalize and convert to dispatcher ops
    gm = functorch.make_fx(functionalize(gm, remove="mutations_and_views"))(*example_inputs)
    gm.graph.print_tabular()
    return gm

@dynamo.optimize(my_backend)
def f(x):
    x[0].relu_()
    return x

m1 = -torch.ones(2, 5)
result = f(m1)
print("m1", m1)
print("result", result)
print("m1 is result:", m1 is result)

opcode         name            target                       args                      kwargs
-------------  --------------  ---------------------------  ------------------------  --------
placeholder    arg0_1          arg0_1                       ()                        {}
call_function  select_copy     aten.select_copy.int         (arg0_1, 0, 0)            {}
call_function  relu            aten.relu.default            (select_copy,)            {}
call_function  select_scatter  aten.select_scatter.default  (arg0_1, relu, 0, 0)      {}
call_function  copy_           aten.copy_.default           (arg0_1, select_scatter)  {}
output         output          output                       ((select_scatter,),)      {}
m1 tensor([[ 0.,  0.,  0.,  0.,  0.],
        [-1., -1., -1., -1., -1.]])
result tensor([[ 0.,  0.,  0.,  0.,  0.],
        [-1., -1., -1., -1., -1.]])
m1 is result: False

This functionalized output is what I want (except that m1 is result is false, so there is a soundness issue :/ ).

Versions

PyTorch version: 2.0.0.dev20221211+cpu
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: Debian GNU/Linux rodete (x86_64)
GCC version: (Debian 12.2.0-3) 12.2.0
Clang version: 14.0.6-2
CMake version: version 3.25.0
Libc version: glibc-2.35

Python version: 3.10.8 (main, Nov 3 2022, 15:17:13) [GCC 12.2.0] (64-bit runtime)
Python platform: Linux-5.19.11-1rodete1-amd64-x86_64-with-glibc2.35
Is CUDA available: False
CUDA runtime version: No CUDA
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==1.24.0rc2
[pip3] torch==2.0.0.dev20221211+cpu
[pip3] torchvision==0.15.0.dev20221211+cpu

cc @ezyang @gchanan @zou3519 @kadeng @msaroufim @bdhirsh @anijain2305 @chauhang @soumith @wconstab @ngimel

[ezyang]

Known problem, if no inputs require grad it goes a different path that doesn't functionalize. Need to fix this