Repository: pytorch
Issue: Get `aot_autograd`'ed graph without `torch.compile` and freeze constants without Inductor context
Labels: triaged,oncall: pt2,oncall: export,module: aotinductor

### 🚀 The feature, motivation and pitch

I am trying to implement eager mode of PT2E quantization on CPU. Currently, the PT2E quantization on CPU is lowered to Inductor by `torch.compile`. The current work flow is like:
- Quantize the model with `Quantizer`, `prepare_pt2e` and `convert_pt2e`. Now the model graph contains `quantize` and `dequantize`.
- Call `torch.compile` so that patterns like `dequantize - float32_op` and more other patterns are fused and lowered in Inductor.
- Constants are also folded in Inductor

For eager mode, the work flow is designed as
- Quantize the model with `Quantizer`, `prepare_pt2e` and `convert_pt2e`. Now the model graph contains `quantize` and `dequantize`. (same as the Inductor mode)
- Get `aot_autograd`'ed graph without `torch.compile`
- Apply fusion passes in Inductor on the graph manually
- Apply constant folding (freezing) passes in Inductor manually

The problems I met are:
1. I cannot run `aot_autograd` on an ordinary graph module directly. The error is `RuntimeError: Graph output must be a (). This is so that we can avoid pytree processing of the outputs.`. See example code 1.
2. Even though we get the aot_autograd graph somehow, we are not able to do constant folding because all constant params become arguments of forward and the freezing path does not know which are constants. In Inductor, there is a context to keep track of the constants. But I don't have that context if I don't call `torch.compile`. See example code 2.

Is there a way to solve the problems without big changes in the PyTorch code? Thanks.

Example code 1:
```python
import torch
import torchvision
import copy
from torch._export import capture_pre_autograd_graph
from torch.ao.quantization.quantize_pt2e import prepare_pt2e, convert_pt2e
import torch.ao.quantization.quantizer.x86_inductor_quantizer as xiq
from torch.ao.quantization.quantizer.x86_inductor_quantizer import X86InductorQuantizer
from torch._dynamo.backends.common import aot_autograd


def my_compiler(gm, example_inputs):
    print(gm)
    return gm

def pt2e_ptq(m, example_inputs):
    m = m.eval()
    exported_model = torch.export.export_for_training(m, example_inputs).module()
    quantizer = X86InductorQuantizer()
    quantizer.set_global(xiq.get_default_x86_inductor_quantization_config())

    prepared_model = prepare_pt2e(exported_model, quantizer)
    _ = prepared_model(*example_inputs)
    converted_model = convert_pt2e(prepared_model)
    torch.ao.quantization.move_exported_model_to_eval(converted_model)

    with torch.no_grad():
        my_backend = aot_autograd(fw_compiler=my_compiler)
        optimized_model = my_backend(converted_model, example_inputs) # !!! error here !!!

    optimized_model(*example_inputs)


if __name__ == "__main__":
    data = torch.randn(16, 3, 224, 224)
    model_fp = torchvision.models.resnet18(weights=torchvision.models.ResNet18_Weights.DEFAULT)
    pt2e_ptq(copy.deepcopy(model_fp), (data,))
```

Example code 2:
```python
import torch
import torchvision
import copy
from torch._export import capture_pre_autograd_graph
from torch.ao.quantization.quantize_pt2e import prepare_pt2e, convert_pt2e
import torch.ao.quantization.quantizer.x86_inductor_quantizer as xiq
from torch.ao.quantization.quantizer.x86_inductor_quantizer import X86InductorQuantizer
from torch._dynamo.backends.common import aot_autograd
from functorch.compile import make_boxed_func
from torch._inductor.fx_passes.freezing_patterns import freezing_passes


def my_compiler(gm, example_inputs):
    freezing_passes(gm, example_inputs) # !!! freezing passes are applied but constant_fold passes are not !!!
    print(gm)
    return make_boxed_func(gm.forward)

my_backend = aot_autograd(fw_compiler=my_compiler)


def pt2e_ptq(m, example_inputs):
    m = m.eval()
    exported_model = torch.export.export_for_training(m, example_inputs).module()
    quantizer = X86InductorQuantizer()
    quantizer.set_global(xiq.get_default_x86_inductor_quantization_config())
    prepared_model = prepare_pt2e(exported_model, quantizer)
    _ = prepared_model(*example_inputs)
    converted_model = convert_pt2e(prepared_model)
    torch.ao.quantization.move_exported_model_to_eval(converted_model)

    with torch.no_grad():
        optimized_model = torch.compile(converted_model, backend=my_backend)

    _ = optimized_model(*example_inputs)

if __name__ == "__main__":

    data = torch.randn(16, 3, 224, 224)
    model_fp = torchvision.models.resnet18(weights=torchvision.models.ResNet18_Weights.DEFAULT)
    pt2e_ptq(copy.deepcopy(model_fp), (data,))
```

### Alternatives

_No response_

### Additional context

_No response_

cc @ezyang @chauhang @penguinwu @avikchaudhuri @gmagogsfm @zhxchen17 @tugsbayasgalan @angelayi @suo @ydwu4 @desertfire @chenyang78