[Mobile GPU] Ban mutations in JIT passes (#56070)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/56070
**Summary**
Currently, we're returning copies instead of alias on mobile GPU (Metal/Vulkan). As suggested by ailzhang , we could use the JIT pass - `RemoveTensorMutation` to ban mutations ahead of time. I've tested two scenarios as shown below. They both work fine on mobile.
- view
```
class Model (torch.nn.Module):
def forward(self, x):
y = x.view(-1)
z = torch.tensor(2.0).float()
y.add_(z)
return x
m = Model()
x = torch.rand(2, 3)
y = m(x)
```
- transpose
```
class Model (torch.nn.Module):
def forward(self, x):
y = x.transpose(1, 2)
z = torch.tensor(2.0).float()
x.add_(z)
return y
m = Model()
x = torch.rand(1, 2, 3)
y = m(x)
```
As we're adding more ops, we should add more tests to cover all the alias ops - https://github.com/pytorch/pytorch/blob/master/tools/autograd/gen_inplace_or_view_type.py#L31-L80
**Next step**
Synced offline with eellison. Since mutation removal is also being used in ONNX, Static runtime, some jit optimizations, Torch -> TVM, etc, instead of inventing something new, we would continue to make it better in cases where it fails.
Although this JIT pass could work for most of the mobile models, there are cases that it can't cover. What we're going to do next is to implement stub ops for GPU models to let them run on server side, such that users can compare results to see if there is any discrepancy.
ghstack-source-id: 126802123
Test Plan:
- Sandcastle
- CircleCI
Reviewed By: raziel
Differential Revision: D27692683
fbshipit-source-id: 9d1be8a6c0a276032b1907807a54fbe2afd882f9
