Arm backend: Keep internal shared qspec with uint8 IO (#20978)

cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils
@Sebastian-Larsson @robell @rascani

Author: bdemirb

backends/arm/quantizer/arm_quantizer_utils.py

@@ -598,11 +598,14 @@ def _model_has_uint8_io(self, model: torch.fx.GraphModule) -> bool:
             self._is_uint8_quantized_io_boundary(node) for node in model.graph.nodes
         )
 
-    def _get_shared_clique(self, root_node: Node) -> tuple[set[Node], list[Any], bool]:
+    def _get_shared_clique(
+        self, root_node: Node
+    ) -> tuple[set[Node], list[Any], bool, bool]:
         shared_nodes = set()
         bfs_queue = [root_node]
         adjacent_qspecs: list[Any] = []
         touches_quantized_io = False
+        touches_uint8_quantized_io = False
 
         while bfs_queue:
             node = bfs_queue.pop(0)
@@ -612,13 +615,24 @@ def _get_shared_clique(self, root_node: Node) -> tuple[set[Node], list[Any], boo
                 self._maybe_enqueue_shared_node(input_node, shared_nodes, bfs_queue)
                 self._append_output_qspec(input_node, adjacent_qspecs)
                 touches_quantized_io |= self._is_quantized_io_boundary(input_node)
+                touches_uint8_quantized_io |= self._is_uint8_quantized_io_boundary(
+                    input_node
+                )
 
             for output_node in node.users.keys():
                 self._maybe_enqueue_shared_node(output_node, shared_nodes, bfs_queue)
                 self._append_input_qspec(output_node, node, adjacent_qspecs)
                 touches_quantized_io |= self._is_quantized_io_boundary(output_node)
+                touches_uint8_quantized_io |= self._is_uint8_quantized_io_boundary(
+                    output_node
+                )
 
-        return shared_nodes, adjacent_qspecs, touches_quantized_io
+        return (
+            shared_nodes,
+            adjacent_qspecs,
+            touches_quantized_io,
+            touches_uint8_quantized_io,
+        )
 
     def _should_skip_while_shared_qspec(self, node: Node) -> bool:
         return node.target == torch.ops.higher_order.while_loop and bool(
@@ -673,9 +687,12 @@ def _annotate_shared_cluster(self, root_node: Node) -> None:
             )
             return
 
-        shared_nodes, adjacent_qspecs, touches_quantized_io = self._get_shared_clique(
-            root_node
-        )
+        (
+            shared_nodes,
+            adjacent_qspecs,
+            touches_quantized_io,
+            touches_uint8_quantized_io,
+        ) = self._get_shared_clique(root_node)
 
         # If there is no neighbor qspec to propagate but the cluster sits on the
         # quantized I/O boundary (e.g. a state-passthrough cat whose only neighbors
@@ -695,6 +712,15 @@ def _annotate_shared_cluster(self, root_node: Node) -> None:
         node_order = {node: index for index, node in enumerate(root_node.graph.nodes)}
         ordered_nodes = sorted(shared_nodes, key=lambda node: node_order.get(node, 0))
 
+        if touches_uint8_quantized_io and any(
+            node.target in self._UINT8_IO_BRIDGE_OPS for node in shared_nodes
+        ):
+            self.report_reject(
+                ordered_nodes,
+                "Shared-qspec bridge cluster touches uint8 model IO.",
+            )
+            return
+
         if self._annotate_while_with_additional_inputs(root_node, adjacent_qspecs):
             return
 
@@ -734,20 +760,9 @@ def _annotate_shared_cluster(self, root_node: Node) -> None:
         return
 
     def annotate(self, model: torch.fx.GraphModule) -> None:  # type: ignore[override]
-        targets = self.targets
-        if self._model_has_uint8_io(model):
-            targets = [
-                target for target in targets if target not in self._UINT8_IO_BRIDGE_OPS
-            ]
-
-        original_targets = self.targets
-        self.targets = targets
-        try:
-            for node in model.graph.nodes:
-                if node.target in self.targets and not self._is_annotated(node):
-                    self._annotate_shared_cluster(node)
-        finally:
-            self.targets = original_targets
+        for node in model.graph.nodes:
+            if node.target in self.targets and not self._is_annotated(node):
+                self._annotate_shared_cluster(node)
 
     def validate(self, model: torch.fx.GraphModule) -> bool:  # type: ignore[override]
         return True

backends/arm/test/quantizer/test_uint8_io_quantization.py

@@ -50,6 +50,19 @@ def forward(self, img0: torch.Tensor, img1: torch.Tensor) -> torch.Tensor:
         return torch.clone(merged)
 
 
+class InternalCatBetweenConvs(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv0 = torch.nn.Conv2d(3, 4, 1)
+        self.conv1 = torch.nn.Conv2d(3, 4, 1)
+        self.conv2 = torch.nn.Conv2d(8, 4, 1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x0 = self.conv0(x)
+        x1 = self.conv1(x)
+        return self.conv2(torch.cat([x0, x1], dim=1))
+
+
 def _get_observer_scale(prepared, observer_node_name: str) -> float:
     observer = prepared.get_submodule(observer_node_name)
     scale, _ = observer.calculate_qparams()
@@ -129,8 +142,36 @@ def test_cat_does_not_bridge_shared_qspec_clusters():
     graph_nodes = {node.name: node for node in prepared.graph.nodes}
     img0_observer = next(iter(graph_nodes["img0"].users))
     img1_observer = next(iter(graph_nodes["img1"].users))
-    final_cat_observer = next(iter(graph_nodes["cat_1"].users))
+    clone_observer = next(iter(graph_nodes["clone"].users))
 
     assert _get_observer_scale(prepared, img0_observer.target) < 0.01
     assert _get_observer_scale(prepared, img1_observer.target) < 0.01
-    assert _get_observer_scale(prepared, final_cat_observer.target) > 0.05
+    assert Q_ANNOTATION_KEY not in graph_nodes["cat_1"].meta
+    assert _get_observer_scale(prepared, clone_observer.target) > 0.05
+
+
+def test_internal_cat_still_shares_qspec_with_uint8_io():
+    """Regression: preserved uint8 model IO must not disable shared-qspec
+    annotation for internal cats between quantized operators.
+    """
+    model = InternalCatBetweenConvs().eval()
+    test_data = (torch.rand(1, 3, 8, 8),)
+    compile_spec = common.get_tosa_compile_spec("TOSA-1.0+INT")
+
+    tosa_quantizer = TOSAQuantizer(compile_spec, use_composable_quantizer=True)
+    tosa_quantizer.set_global(get_symmetric_quantization_config())
+    tosa_quantizer.set_io(get_uint8_io_quantization_config())
+
+    exported = torch.export.export(model, test_data, strict=True)
+    prepared = prepare_pt2e(exported.module(), tosa_quantizer)
+
+    cat_nodes = [
+        n
+        for n in prepared.graph.nodes
+        if n.op == "call_function" and n.target == torch.ops.aten.cat.default
+    ]
+    assert len(cat_nodes) == 1, f"Expected 1 cat node, got {len(cat_nodes)}"
+    cat_node = cat_nodes[0]
+
+    assert Q_ANNOTATION_KEY in cat_node.meta
+    assert cat_node.meta[Q_ANNOTATION_KEY].output_qspec is not None