Repository: pytorch
Issue: [hierarchical compilation] A way to designate a portion of torch.compile as a noinline block that is compiled/guarded separately, but less disruptive than a graph break (e.g., for loops)
Labels: feature,triaged,oncall: pt2,module: dynamo,module: startup-tracing-compile

### 🐛 Describe the bug

There are few motivating uses cases for this:

* Compile time, previously filed as https://github.com/pytorch/pytorch/issues/102839
* Varying the compile-time options for different regions of the Dynamo graph

The essential model here is that I may have some subcomponent of my torch.compile program, which I would like to compile separately from the rest for some reason. In the compile time use case, maybe this is a loop body that I don't want to unroll 20 times, I want to compile only once and reuse it all those times. In the varying compile-time use case, I might have a regular model that I want to have regular dynamic shapes support, but in this particular region, I want my code all specialized with static shapes.

Alternatives:

* The ordinary tool for breaking up a torch.compile program is introducing a graph break. However, because we don't graph breaking inside inlined functions https://github.com/pytorch/pytorch/issues/111003 a graph break can end up being pretty disruptive to the rest of the graph structure. Also you might consider having the "whole" graph for export reasons, so you don't want a "true" graph break, you want something a little weaker.
* For the loop use case, you could potentially alternately just define a higher order op that supports looping; typically the higher order op would support any Python mutation on the inside but maybe that's fine.

So, what exactly would this do?

* The inner region should have its own guards. In fact, it needs an entirely separate Dynamo instance. For loop bodies, this is important because there may be loop carried variables that get updated, and we need to reify the side effects (e.g., the normal graph break way). For the vary compile time options, if you have a static segment inside an otherwise dynamic program, you need to do a guard dispatch to go to the correct static program.
* The outer Dynamo instance should somehow "incorporate" the inner Dynamo instance in some way. For example, it can trace directly through the modified bytecode produced by the inner Dynamo instance to apply the loop. But instead of tracing into the compiled function, it just emits a call directly to it (and the compiled function is assumed to be well behaved in some sense)
* The outer Dynamo instance uses a partitioner, similar to optimize ddp partitioner, to separate out the already compiled products so we don't try to send them to AOTAutograd

So for example, suppose you have:

```
x = f(y)
for i in range(3):
  x = g(x)
return x
```

And suppose g is marked to be a block. Then in normal Dynamo tracing we trace up until the g call, and then we trigger an inner Dynamo instance which produces a compiled g bytecode which we trace directly into the graph. The outer graph finally looks like:

```
x0 = f(y)  # assume f is inlined here
x1 = compiled_g_no_residual_bytecode(x0)
x2 = compiled_g_no_residual_bytecode(x1)
x3 = compiled_g_no_residual_bytecode(x2)
return x3
```

And then we partition the graph into four segments and then compile them individually (the compiled segments don't need more compiling, they just pass through).

### Versions

main

cc @anijain2305 @chauhang @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @oulgen @jamesjwu @aorenste @laithsakka @msaroufim @bdhirsh @zou3519 @aakhundov