Repository: pytorch
Issue: `min_cut_rematerialization_partition` is not optimal memory usage.
Labels: triaged,topic: performance,oncall: pt2,module: inductor

Test code: 
```python
import torch
from functorch.compile import aot_function, make_boxed_func, min_cut_rematerialization_partition

def fn(x, y, z, a, b, c):
    return (a * c) * x + (b) * y + ((a * b) * c) * z

def run_func(func, *inputs):
    res = func(*inputs)
    loss = res.sum()
    loss.backward()

def compiler_fn(fx_module: torch.fx.GraphModule, _):
    print(fx_module.code)
    return make_boxed_func(fx_module.forward)

x, y, z, a, b, c = [torch.randn(2, 4, requires_grad=i < 3, device="cuda") for i in range(6)]

aot_mincut_fn = aot_function(fn, fw_compiler=compiler_fn,
    bw_compiler=compiler_fn,
    partition_fn=min_cut_rematerialization_partition,
    decompositions=None)
run_func(aot_mincut_fn, x, y, z, a, b, c)
```

Output:
```
def forward(self, primals_1, primals_2, primals_3, primals_4, primals_5, primals_6):
    mul = torch.ops.aten.mul.Tensor(primals_4, primals_6)
    mul_1 = torch.ops.aten.mul.Tensor(mul, primals_1);  mul = primals_1 = None
    mul_2 = torch.ops.aten.mul.Tensor(primals_5, primals_2);  primals_2 = None
    add = torch.ops.aten.add.Tensor(mul_1, mul_2);  mul_1 = mul_2 = None
    mul_3 = torch.ops.aten.mul.Tensor(primals_4, primals_5)
    mul_4 = torch.ops.aten.mul.Tensor(mul_3, primals_6);  mul_3 = None
    mul_5 = torch.ops.aten.mul.Tensor(mul_4, primals_3);  mul_4 = primals_3 = None
    add_1 = torch.ops.aten.add.Tensor(add, mul_5);  add = mul_5 = None
    return [add_1, primals_4, primals_5, primals_6]
    



def forward(self, primals_4, primals_5, primals_6, tangents_1):
    mul = torch.ops.aten.mul.Tensor(primals_4, primals_6)
    mul_3 = torch.ops.aten.mul.Tensor(primals_4, primals_5);  primals_4 = None
    mul_4 = torch.ops.aten.mul.Tensor(mul_3, primals_6);  mul_3 = primals_6 = None
    mul_6 = torch.ops.aten.mul.Tensor(tangents_1, mul_4);  mul_4 = None
    mul_7 = torch.ops.aten.mul.Tensor(tangents_1, primals_5);  primals_5 = None
    mul_8 = torch.ops.aten.mul.Tensor(tangents_1, mul);  tangents_1 = mul = None
    return [mul_8, mul_7, mul_6, None, None, None]

```

Explain:
We assume `x`, `y`, `z` require grad, but `a`, `b`, `c` do not require grad. `f = ac x + b y + abc z`. ∂f/∂x = ac, ∂f/∂y = b, ∂f/∂z = abc. Optimally, we only need to store `ac` and `b`, because we can compute `abc` with `ac · b`, it's just Hadamard product. But it output `a`, `b`, `c`, which consumes more memory. If let `g = ac x + b y + acb z`, it could return the optimal answer. So considering commutative law and combined law could get better result. That is to say, we need some meta information of the operation. Otherwise when building max-flow-min-cut graph, if `A = op(B, C, ...)`, where `A`, `B`, `C`, `...` is the node in compute graph, then we also link them. Otherwise let smart compiler reorder the code, let `abc` to `acb` so that current solution works. However, it takes more complexity during initialization.

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