Repository: pytorch
Issue: [RFC] Optional Modular Representation for FX Graph from `torch.compile()`
Labels: feature,triaged,oncall: pt2

### 🚀 The feature, motivation and pitch

Currently, when using TorchDynamo to generate an FX graph, it inlines all modules and function calls. Such an inlining can introduce extra burden to some custom compilers since the same layer needs to be recompiled multiple times. Ideally, for the same layer that's being called multiple times, we only want to compile it once. There are majorly two solutions. 

First, which is what we are doing right now, is to reverse-engineer the layer information from the flattend FX graph. However, this only works well for well-known models since we know how a layer looks like. For other custom models, it would be hard to tell (although we can still retrieve some information from the module_stack metadata).

Second, which is what we are proposing, is to enable modular representation inside FX graph (i.e., we do not inline modules). However, this approach could introduce extra effort for compilers to perform optimizations. Thus, we propose to make it an optional feature. 

Following is an example of how this modular representation looks like.

```python
# input program
class MyModule(torch.nn.Module):

    class MySubModule(torch.nn.Module):
        def __init__(self):
            super().__init__()

        def forward(self, x):
            return x+2

    def __init__(self):
        super().__init__()
        self.linear = torch.nn.Linear(4, 5)
        self.submod = self.MySubModule()

    def forward(self, x):
        return self.linear(self.submod(x))

# generated code
class GraphModule(torch.nn.Module):
    class SubGraphModule(torch.nn.Module):
         def forward(self, L_x_ : torch.Tensor):
            l_x_ = L_x_                                                                                                       
            add = l_x_ + 2;  l_x_ = None                                                                                      
            return (add,)
    def forward(self, L_x_ : torch.Tensor):
        l_x_ = L_x_ 
        l__self___submod = self.SubGraphModule()
        l__self___submod_1 =  l__self___submod(l_x_);  l_x_ = None                                                                                                                                                                                            
        l__self___linear = self.L__self___linear(l__self___submod_1);  l__self___submod_1 = None                  
        return (l__self___linear,)   
```

To enable this representation, we will need to following
1. Introduce two new FX graph nodes: one for the submodule itself, the other one for instatiating and calling the submodule.
2. Introduce a new FX pass that inlines the submodule.

Link to prototype: https://github.com/seanlatias/pytorch/tree/modular

cc: @SherlockNoMad @wconstab 

### Alternatives

_No response_

### Additional context

Our final goal is to connect this flow with torch-xla and generate StableHLO code. In other words, we want to keep this modular representation at the MLIR level and inline there if needed. 

Another use case that we have in mind for this feature, is to generate a unified training graph from `torch.compile()`, and utilize this modular representation to mark the boundary between forward, backward, and optimizer. With this approach, the compiler would have more flexibility on whether to inline or not.

cc @ezyang @msaroufim @bdhirsh @anijain2305 @zou3519 @chauhang @wconstab