Repository: pytorch
Issue: Improve shape padding in training.
Labels: module: performance,feature,triaged,oncall: pt2,module: inductor

### 🚀 The feature, motivation and pitch

Our shape padding currently doesn't really do that well in training. In training for code like
```
def f(x):
    x = x + 1
    return torch.mm(x, x) + 1
```

we get a graph that looks like this in the backwards.
```
   def forward(self, full_default_8: f16[4095, 1], full_default_9: f16[1, 4095], full_default_10: f16[4096, 1], full_default_11: f16[1, 4096], permute: f16[4095, 4095], tangents_1: f16[4095, 4095]):
        # No stacktrace found for following nodes
        cat_default_4: f16[4095, 4096] = torch.ops.aten.cat.default([permute, full_default_8], 1)
        cat_default_5: f16[4096, 4095] = torch.ops.aten.cat.default([tangents_1, full_default_9])
        cat_default_6: f16[4096, 4096] = torch.ops.aten.cat.default([cat_default_5, full_default_10], 1);  cat_default_5 = None
        cat_default_7: f16[4096, 4096] = torch.ops.aten.cat.default([cat_default_4, full_default_11]);  cat_default_4 = None
        mm_default_5: f16[4096, 4096] = torch.ops.aten.mm.default(cat_default_7, cat_default_6);  cat_default_7 = cat_default_6 = None
        slice_tensor_6: f16[4095, 4096] = torch.ops.aten.slice.Tensor(mm_default_5, 0, 0, -1);  mm_default_5 = None
        slice_tensor_7: f16[4095, 4096] = torch.ops.aten.slice.Tensor(slice_tensor_6, 1, 0, 9223372036854775807);  slice_tensor_6 = None
        slice_tensor_4: f16[4095, 4096] = torch.ops.aten.slice.Tensor(slice_tensor_7, 0, 0, 9223372036854775807);  slice_tensor_7 = None
        slice_tensor_5: f16[4095, 4095] = torch.ops.aten.slice.Tensor(slice_tensor_4, 1, 0, -1);  slice_tensor_4 = None
        cat_default: f16[4095, 4096] = torch.ops.aten.cat.default([tangents_1, full_default_8], 1);  tangents_1 = full_default_8 = None
        cat_default_1: f16[4096, 4095] = torch.ops.aten.cat.default([permute, full_default_9]);  permute = full_default_9 = None
        cat_default_2: f16[4096, 4096] = torch.ops.aten.cat.default([cat_default_1, full_default_10], 1);  cat_default_1 = full_default_10 = None
        cat_default_3: f16[4096, 4096] = torch.ops.aten.cat.default([cat_default, full_default_11]);  cat_default = full_default_11 = None
        mm_default_2: f16[4096, 4096] = torch.ops.aten.mm.default(cat_default_3, cat_default_2);  cat_default_3 = cat_default_2 = None
        slice_tensor_2: f16[4095, 4096] = torch.ops.aten.slice.Tensor(mm_default_2, 0, 0, -1);  mm_default_2 = None
        slice_tensor_3: f16[4095, 4096] = torch.ops.aten.slice.Tensor(slice_tensor_2, 1, 0, 9223372036854775807);  slice_tensor_2 = None
        slice_tensor: f16[4095, 4096] = torch.ops.aten.slice.Tensor(slice_tensor_3, 0, 0, 9223372036854775807);  slice_tensor_3 = None
        slice_tensor_1: f16[4095, 4095] = torch.ops.aten.slice.Tensor(slice_tensor, 1, 0, -1);  slice_tensor = None
        
        # File: ../t.py:8, code: return torch.mm(x, x) + 1
        add_2: f16[4095, 4095] = torch.ops.aten.add.Tensor(slice_tensor_1, slice_tensor_5);  slice_tensor_1 = slice_tensor_5 = None
        return [add_2]
```

The problem is that the joint graph looks like
```
add = ...
mm: f16[4095, 4095] = torch.ops.aten.mm.default(add, add)

permute: f16[4095, 4095] = torch.ops.aten.permute.default(add, [1, 0])
mm_1: f16[4095, 4095] = torch.ops.aten.mm.default(permute, tangents_1);  permute = None
permute_1: f16[4095, 4095] = torch.ops.aten.permute.default(add, [1, 0]);  add = None
mm_2: f16[4095, 4095] = torch.ops.aten.mm.default(tangents_1, permute_1);  tangents_1 = permute_1 = None
```
So even after we say, pad `add` in the forwards pass, the way the tensor is used in the backwards pass depends on `permute(add)`, which has *not* been padded.

Ideally we should be able to avoid this, either by having a smarter joint graph pattern, or by doing our matmul padding on the pre-autograd graph? Not totally sure.

### Alternatives

_No response_

### Additional context

_No response_

cc @ezyang @msaroufim @wconstab @bdhirsh @anijain2305 @zou3519 @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @aakhundov @ColinPeppler @Xia-Weiwen @ngimel