Update base for Update on "Support export program in intermediate numeric discrepancy detector"

This diff enables intermediate numeric discrepancy detector to leverage export program as label.

More specific, if user creates etrecord with exported program, and the exported program is one of the exported programs in the export flow, then our numeric discrepancy detector will use it as label. Otherwise, we will continue use edge dialect graph as label.

Differential Revision: [D78298935](https://our.internmc.facebook.com/intern/diff/D78298935/)

[ghstack-poisoned]

Author: Gasoonjia

.ci/docker/ci_commit_pins/pytorch.txt

@@ -1 +1 @@
-ab43fe4bdf5ccd82897f0e982c451a0127bd175e
+2dccff7dcf56b0d168ebfd7ca08bdeca37273c56

backends/arm/arm_backend.py

@@ -128,7 +128,7 @@ def ethosu_compile_spec(
         self.compiler_flags.append("--output-format=raw")
         self.compiler_flags.append("--debug-force-regor")
 
-        base_tosa_version = "TOSA-0.80+BI"
+        base_tosa_version = "TOSA-1.0+INT"
         if "u55" in target:
             # Add the Ethos-U55 extension marker
             base_tosa_version += "+u55"

backends/arm/quantizer/arm_quantizer_utils.py

@@ -18,22 +18,21 @@
 from torch.fx import GraphModule, Node
 
 from torchao.quantization.pt2e.quantizer import QuantizationAnnotation
+from torchao.quantization.pt2e.quantizer.quantizer import Q_ANNOTATION_KEY
 
 
 def is_annotated(node: Node) -> bool:
     """Given a node return whether the node is annotated."""
     return (
-        "quantization_annotation" in node.meta
-        and cast(
-            QuantizationAnnotation, node.meta["quantization_annotation"]
-        )._annotated
+        Q_ANNOTATION_KEY in node.meta
+        and cast(QuantizationAnnotation, node.meta[Q_ANNOTATION_KEY])._annotated
     )
 
 
 def is_output_annotated(node: Node) -> bool:
     """Given a node, return whether the output of the node is annotated."""
-    if "quantization_annotation" in node.meta:
-        annotation = cast(QuantizationAnnotation, node.meta["quantization_annotation"])
+    if Q_ANNOTATION_KEY in node.meta:
+        annotation = cast(QuantizationAnnotation, node.meta[Q_ANNOTATION_KEY])
         return annotation._annotated and annotation.output_qspec is not None
     else:
         return False
@@ -43,9 +42,9 @@ def mark_node_as_annotated(node: Node) -> None:
     """Marks node as annotated. If needed, an empty  QuantizationAnnotation is added
     to the quantization_annotation node meta entry.
     """
-    if "quantization_annotation" not in node.meta:
-        node.meta["quantization_annotation"] = QuantizationAnnotation()
-    node.meta["quantization_annotation"]._annotated = True
+    if Q_ANNOTATION_KEY not in node.meta:
+        node.meta[Q_ANNOTATION_KEY] = QuantizationAnnotation()
+    node.meta[Q_ANNOTATION_KEY]._annotated = True
 
 
 def is_ok_for_quantization(node: Node, gm: GraphModule):

backends/arm/test/test_arm_baremetal.sh

@@ -228,7 +228,7 @@ test_models_ethos-u85() { # End to End model tests using model_test.py
     python3 backends/arm/test/test_model.py --test_output=arm_test/test_model --target=ethos-u85-256 --model=mv2 --extra_flags="-DET_ATOL=2.00 -DET_RTOL=2.00"
     python3 backends/arm/test/test_model.py --test_output=arm_test/test_model --target=ethos-u85-512 --model=mv3 --extra_flags="-DET_ATOL=5.00 -DET_RTOL=5.00"
     python3 backends/arm/test/test_model.py --test_output=arm_test/test_model --target=ethos-u85-128 --model=lstm --extra_flags="-DET_ATOL=0.03 -DET_RTOL=0.03"
-    python3 backends/arm/test/test_model.py --test_output=arm_test/test_model --target=ethos-u85-128 --model=w2l --extra_flags="-DET_ATOL=0.01 -DET_RTOL=0.01"
+    #python3 backends/arm/test/test_model.py --test_output=arm_test/test_model --target=ethos-u85-128 --model=w2l --extra_flags="-DET_ATOL=0.01 -DET_RTOL=0.01"  # Takes long time to run
     python3 backends/arm/test/test_model.py --test_output=arm_test/test_model --target=ethos-u85-256 --model=ic4 --extra_flags="-DET_ATOL=0.8 -DET_RTOL=0.8" --timeout=2400
     python3 backends/arm/test/test_model.py --test_output=arm_test/test_model --target=ethos-u85-128 --model=resnet18 --extra_flags="-DET_ATOL=0.2 -DET_RTOL=0.2"
     python3 backends/arm/test/test_model.py --test_output=arm_test/test_model --target=ethos-u85-128 --model=resnet50 --extra_flags="-DET_ATOL=0.2 -DET_RTOL=0.2"

backends/cadence/aot/compiler.py

@@ -228,6 +228,9 @@ def _lower_ep_to_edge(
     """
     Lower an ExportedProgram to an EdgeProgramManager (in edge IR).
     """
+    # Apply passes which transform the ExportedProgram before it gets lowered to edge.
+    expo_program = apply_torch_ops_passes(expo_program)
+
     # Call to_edge to convert the graph to edge IR.
     # Note: dim_order is skipped (https://github.com/pytorch/executorch/issues/3704)
     edge_prog_manager = to_edge(
@@ -263,9 +266,6 @@ def export_to_edge(
     # Export the model into an ExportedProgram.
     expo_program = trace(model, inputs)
 
-    # Apply passes which transform the ExportedProgram before it gets lowered to edge.
-    expo_program = apply_torch_ops_passes(expo_program)
-
     # Lower the model to edge IR.
     edge_prog_manager = _lower_ep_to_edge(
         expo_program, dump_graphs, constant_methods, core_aten_exceptions

backends/cadence/aot/ops_registrations.py

@@ -276,6 +276,14 @@
     "requantize.per_tensor_out(Tensor input, float in_scale, int in_zero_point, float out_scale, "
     "int out_zero_point, ScalarType out_dtype, *, Tensor(a!) out) -> Tensor(a!)"
 )
+lib.define(
+    "roi_align_box_processor.out(Tensor rois, int output_size_h, int output_size_w, "
+    "int sampling_ratio, bool aligned, *, Tensor(a!) out) -> Tensor(a!)"
+)
+lib.define(
+    "roi_align_box_processor(Tensor rois, int output_size_h, int output_size_w, "
+    "int sampling_ratio, bool aligned) -> (Tensor out)"
+)
 
 # Custom ops with aten namespace. Need to specify the lib var as FRAGMENT type as aten library is already defined
 aten_lib = Library("aten", "FRAGMENT")
@@ -1038,3 +1046,14 @@ def idma_store_impl(
     channel: int = 0,
 ) -> torch.Tensor:
     return copy_idma_copy_impl(src, task_num, channel)
+
+
+@register_fake("cadence::roi_align_box_processor")
+def roi_align_box_processor_meta(
+    rois: torch.Tensor,
+    output_size_h: int,
+    output_size_w: int,
+    sampling_ratio: int,
+    aligned: bool,
+) -> torch.Tensor:
+    return rois.new_empty((rois.shape[0], 80), dtype=torch.uint8)

backends/cadence/aot/quantizer/quantizer.py

@@ -42,6 +42,7 @@
     QuantizationSpec,
     Quantizer,
 )
+from torchao.quantization.pt2e.quantizer.quantizer import Q_ANNOTATION_KEY
 
 
 act_qspec_asym8s = QuantizationSpec(
@@ -127,7 +128,7 @@ def annotate(self, model: torch.fx.GraphModule) -> torch.fx.GraphModule:
 
             for output, *custom_spec in anchors.output:
                 # pyre-ignore[16]: no attribute
-                output.meta["quantization_annotation"] = QuantizationAnnotation(
+                output.meta[Q_ANNOTATION_KEY] = QuantizationAnnotation(
                     # pyre-ignore[6]: incompatible parameter type
                     output_qspec=(custom_spec[0] if custom_spec else output_act_qspec),
                     _annotated=True,
@@ -143,7 +144,7 @@ def annotate_inputs(
                 for node, idx, *custom_spec in inputs:
                     # pyre-ignore[16]: no attribute
                     annotation = node.meta.get(
-                        "quantization_annotation",
+                        Q_ANNOTATION_KEY,
                         QuantizationAnnotation(_annotated=True),
                     )
                     arg = (
@@ -157,21 +158,21 @@ def annotate_inputs(
                         custom_spec[0] if custom_spec else spec
                     )
                     # pyre-ignore[16]: no attribute
-                    node.meta["quantization_annotation"] = annotation
+                    node.meta[Q_ANNOTATION_KEY] = annotation
 
             def annotate_weights_or_biases(
                 weights_or_biases: List[Tuple[fx.Node, int]],
                 spec: Optional[QuantizationSpec],
             ) -> None:
                 for node, idx, *custom_spec in weights_or_biases:
                     annotation = node.meta.get(
-                        "quantization_annotation",
+                        Q_ANNOTATION_KEY,
                         QuantizationAnnotation(_annotated=True),
                     )
                     annotation.input_qspec_map[node.args[idx]] = (
                         custom_spec[0] if custom_spec else spec
                     )
-                    node.meta["quantization_annotation"] = annotation
+                    node.meta[Q_ANNOTATION_KEY] = annotation
 
             # pyre-ignore[6]: incompatible parameter type
             annotate_inputs(anchors.inputs, input_act_qspec)

backends/cadence/aot/quantizer/utils.py

@@ -21,6 +21,7 @@
     SourcePartition,
 )
 from torchao.quantization.pt2e import ObserverOrFakeQuantize
+from torchao.quantization.pt2e.quantizer.quantizer import Q_ANNOTATION_KEY
 
 
 def quantize_tensor_multiplier(
@@ -88,8 +89,7 @@ def is_annotated(nodes: List[fx.Node]) -> bool:
     annotated = False
     for node in nodes:
         annotated = annotated or (
-            "quantization_annotation" in node.meta
-            and node.meta["quantization_annotation"]._annotated
+            Q_ANNOTATION_KEY in node.meta and node.meta[Q_ANNOTATION_KEY]._annotated
         )
     return annotated
 

backends/cadence/aot/replace_ops.py

@@ -2328,12 +2328,15 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
 
             # Extract an argument to a separate full op.
             with graph_module.graph.inserting_before(mul_node):
-                full_tensor = graph_module.graph.call_function(
+                full_node = graph_module.graph.call_function(
                     torch.ops.aten.full.default, args=([1], full_arg)
                 )
+                full_node.meta = mul_node.meta
+                full_node.meta["val"] = [1]
                 new_mul_node = graph_module.graph.call_function(
-                    torch.ops.aten.mul.Tensor, args=(x_arg, full_tensor)
+                    torch.ops.aten.mul.Tensor, args=(x_arg, full_node)
                 )
+                new_mul_node.meta = mul_node.meta
             # Replace the old mul with a newly created mul.
             mul_node.replace_all_uses_with(new_mul_node)
             graph_module.graph.erase_node(mul_node)

backends/cortex_m/test/test_replace_quant_nodes.py

@@ -25,6 +25,7 @@
     QuantizationSpec,
     Quantizer,
 )
+from torchao.quantization.pt2e.quantizer.quantizer import Q_ANNOTATION_KEY
 
 
 @dataclass(eq=True, frozen=True)
@@ -67,18 +68,15 @@ def annotate(self, model: GraphModule):
             ]:
                 continue
 
-            if (
-                "quantization_annotation" in node.meta
-                and node.meta["quantization_annotation"]._annotated
-            ):
+            if Q_ANNOTATION_KEY in node.meta and node.meta[Q_ANNOTATION_KEY]._annotated:
                 continue
 
             input_qspec_map = {
                 node.args[0]: config.input_activation,
                 node.args[1]: config.input_activation,
             }
 
-            node.meta["quantization_annotation"] = QuantizationAnnotation(
+            node.meta[Q_ANNOTATION_KEY] = QuantizationAnnotation(
                 input_qspec_map=input_qspec_map,
                 output_qspec=config.output_activation,
                 _annotated=True,