Update base for Update on "make operator name consistent before and after serde"

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

[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,