Repository: pytorch
Issue: [JIT] Allow user to provide aliasing information on input tensors and model parameters
Labels: triage review,oncall: jit,triaged

## 🚀 Feature
<!-- A clear and concise description of the feature proposal -->

We need to create API that allows user to provide aliasing information on input tensors and model parameters. Graph mode does not own parameters and takes them as input through `prim::GetAttr` and `ClassType<xxx>`. User has full visibility of aliasing information on direct tensor inputs and model parameters. It is possible for user to provide aliasing information when scripting their model.

In default configuration where there's no user-feed aliasing information, `AliasDb` could be conservative and preserve current assumption that all inputs could be aliasing each other. The proposed feature would not change existing behavior.

## Motivation

<!-- Please outline the motivation for the proposal. Is your feature request related to a problem? e.g., I'm always frustrated when [...]. If this is related to another GitHub issue, please link here too -->

It is only rare case where inputs and parameters could be aliasing each other. In such cases the alias relationships should be well-aware to the module creator/user. Hence it is reasonable to expect the module creator/user to provide such information to graph compiler.

Aliasing information on input tensor and model parameters is beneficial to graph optimization. It allows identify aliases on inputs, which are not visible from within the graph. E.g. it could remove many wildcard (`AliasDb` concept of could aliasing anything) in our current `AliasDb` resulted from `prim::GetAttr`, when related alias information is provided for retrieved attributes.
As long as we support in-place operator in JIT, we'll always have input-output tensor (e.g. tensors feed to in-place operator), which would prevent any re-ordering or nodes across it if it is a wildcard.

## Pitch

<!-- A clear and concise description of what you want to happen. -->

Given a `torch.nn.Module` defined as:
```python
class MyModule(torch.nn.Module):
  def __init__(self, inp_c, out_c, inter_c):
    super(MyModule, self).__init__()
    self.fc = torch.nn.modules.Linear(inp_c, int_c)
    self.fc_dup = torch.nn.modules.Linear(int_c, out_c)
    self.fc_dup.weight = self.fc.weight

  # exposed function, aliasing info expected from user
  @torch.jit.export
  def exposed_func(self, tensor1 : torch.Tensor, tensor2 : torch.Tensor) -> torch.Tensor:
    tensor1 += 1e-5
    return tensor2 / tensor1

  # aliasing info expected from user
  def forward(self, inp : torch.Tensor) -> torch.Tensor:
    inp = self.fc(inp)
    inp = self.fc_dup(inp)
    out = inp.relu_()
    # This call instance of `exposed_func` is within the graph,
    #   aliasing info on inputs should be deduced through analysis
    #   instead of passed in by user.
    return self.exposed_func(inp, out)
```
With current scripting API, we could add another field `alias_info` in `torch.jit.script` function to receive aliasing information.
```
script_model = torch.jit.script(MyModule, alias_info = alias_info)
```
where `alias_info` could be 
```
Dict{
    'forward' : 
        [['0.fc.weight', '0.fc_dup.weight'],
        ],
    'exposed_func' : 
        [[]],
}
```

Note:

1. For each entry in the dictionary: key matches the compiled method, to which the aliasing information is provided; values are a list of ValueSet, which each is a group of values that are aliases to each other.
  1.1 We refer to inputs directly using their index;
  1.2 We use scoped name to refer to parameters and parameters in nested `torch.nn.Modules`. The example here uses `0.fc.weight`, which is roughly `input_0 (MyModule) -> attribute["fc"] (Linear) -> attribute["weight"] (Parameter)`
2. We have provided `alias_info` for `exposed_func`. This is only used when we compile `exposed_func` directly, which is exposed through annotation `torch.jit.export`. When `exposed_func` is compiled recursively (nested inside `forward`), aliasing information should be deduced by alias analysis in the graph instead of passed in by the user.
3. We provide `alias_info` for both exposed function `forward` and `exposed_func`. If we have a third exposed function here, for which there's no provided alias_info, `AliasDb` could keep the old assumption that all inputs could be aliasing each other.

cc @suo