Adds a "reinplacing" FX transform, that goes through an FX graph and tries to convert out-of-place op calls into inplace calls whenever possible.
Followups from this PR include:
- Set up torch bench, and run the whole torchbench suite using AOTAutograd + functionalize + rein placing transforms to surface any issues (this is what I'm currently working on). Right now, I have some basic unit tests just to sanity check that the general logic makes sense.
- Add any missing inplace ops. This is mostly the `*_scatter*` ops, e.g. `diagonal_scatter_`, because these ops will commonly show up an FX graph after running functionalization.
The criteria for when you can swap an op `b = a.add(...)` with `a.add_(...)` is:
(1) An inplace variant of the operator with the same schema needs to exist (`aten.add` -> `aten.add_`)
(2) `a` (**or any of its aliases**) can't be used as an input to any other operators later on in the graph
(3) `a` can't be one of the inputs to the entire graph. It also can't be an **alias** of any of the inputs ***
*** One thing to note: (3) means that we can't technically guarantee that we'll get back **all** memory usage that we lost from functionalization. Functionalization converts input mutations into out-of-place calls, and then adds a `copy_()` to the end of the graph to preserve semantics.
I added logic to handle `copy_()` in this PR because it it's a pretty important optimizations in the context of `functionalization()`: any program that performs input mutations will have a `copy_()` in it after running functionalization.
There are some examples in the test file, but I think staring at an example of where re-inplacing is/isn't allowed to run is helpful:
```
// Before functionalization
def foo(a):
tmp1 = a.add_(1)
tmp2 = a.add(2)
// After functionalization
def foo(a)
tmp1 = a.add(1)
tmp2 = a.add(2)
....
a.copy_(tmp1)
// After re-inplacing
def foo(a)
// first add() is safe to re-inplace even though a is a program input,
// because a's data is overwritten later by a copy_()
tmp1 = a.add_(1)
// second add() is NOT safe to re-inplace, because:
// (1) a and tmp1 are aliased. Note that they weren't aliased in the original program,
but they are now that we've done some re-inplacing.
// (2) tmp1 is used as an input later in the program
tmp2 = a.add(2)
....
a.copy_(tmp1)
```
Pull Request resolved: https://github.com/pytorch/pytorch/pull/80897
Approved by: https://github.com/ezyang
