Fix Long Term Quant Testing (#78)

Summary:
Pull Request resolved: #78

Long term Quant seems to be using their custom graph capture before running prepare and convert. Let us mirror this so we are testing with the proper quant flow

https://fb.workplace.com/groups/257735836456307/permalink/545316467698241/

The change here would be that Quantize2 would run on the stage before Export.

Testing wise the stages are as follows

```
export.capture_pre_autograd_graph --> prepare --> convert --> exir.capture()

|--------------------------------------------------------|    |-------------|
                           Quantize2(stage)                    Export(stage)
```

Differential Revision: D48488929

fbshipit-source-id: 229b1783cc2c89222e3e5d19216c1fd645e30e82

Author: mcr229

backends/xnnpack/test/ops/add.py

@@ -75,9 +75,9 @@ def test_add_quantized_pt2e(self):
 
         (
             Tester(add_module, model_inputs)
+            .quantize2()
             .export()
             .check_count({"torch.ops.aten.add.Tensor": 4})
-            .quantize2()
             .check(["torch.ops.quantized_decomposed"])
             .to_edge()
             .check_count({"executorch_exir_dialects_edge__ops_aten_add_Tensor": 4})

backends/xnnpack/test/tester/tester.py

@@ -10,6 +10,7 @@
 from typing import Any, Dict, List, Optional, Tuple
 
 import torch
+import torch._export as export
 from executorch import exir
 from executorch.backends.xnnpack.partition.xnnpack_partitioner import (
     XnnpackFloatingPointPartitioner,
@@ -145,28 +146,28 @@ def __init__(
 
         self.quantizer.set_global(self.quantization_config)
 
-        self.converted_program = None
+        self.converted_graph = None
 
     def run(
-        self, artifact: ExirExportedProgram, inputs: Optional[Tuple[torch.Tensor]]
+        self, artifact: torch.nn.Module, inputs: Optional[Tuple[torch.Tensor]]
     ) -> None:
-        prepared = prepare_pt2e(artifact.exported_program.graph_module, self.quantizer)
+        captured_graph = export.capture_pre_autograd_graph(artifact, inputs)
+        prepared = prepare_pt2e(captured_graph, self.quantizer)
         converted = convert_pt2e(prepared)
-        artifact.exported_program._graph_module = converted
-        self.converted_program = artifact
+        self.converted_graph = converted
 
     @property
-    def artifact(self) -> ExirExportedProgram:
-        return self.converted_program
+    def artifact(self) -> torch.fx.GraphModule:
+        return self.converted_graph
 
     @property
     def graph_module(self) -> str:
-        return self.converted_program.exported_program.graph_module
+        return self.converted_graph
 
 
 @register_stage
 class Export(Stage):
-    def __init__(self, capture_config: Optional[CaptureConfig] = None):
+    def __init__(self, for_quant=False, capture_config: Optional[CaptureConfig] = None):
         self.capture_conf = capture_config or get_xnnpack_capture_config()
         self.exir_exported_program = None
 
@@ -274,12 +275,11 @@ def __init__(
         self.inputs = inputs
         self.stages: Dict[str, Stage] = OrderedDict.fromkeys(list(_stages_.keys()))
         self.pipeline = {
+            self._stage_name(Quantize2): [self._stage_name(Export)],
             self._stage_name(Quantize): [self._stage_name(Export)],
             self._stage_name(Export): [
-                self._stage_name(Quantize2),
                 self._stage_name(ToEdge),
             ],
-            self._stage_name(Quantize2): [self._stage_name(ToEdge)],
             self._stage_name(ToEdge): [self._stage_name(Partition)],
             # TODO Make this Stage optional
             self._stage_name(Partition): [self._stage_name(ToExecutorch)],