Arm backend: Avoid shared qspec bridges for uint8 IO (#20976)

When preserved uint8 model IO is used, SharedQspecQuantizer can merge
image-scale tensors with wider internal activation ranges.

Avoid using cat, concatenate, stack, pixel_shuffle, and slice as
fallback shared-qspec bridge ops for graphs with uint8 IO. This keeps
normal int8 quantisation behaviour unchanged while preserving image
input qparams for uint8 IO flows.

Add a regression test covering image-like uint8 IO joined with a
high-range branch.

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

Signed-off-by: Baris Demir <baris.demir@arm.com>

Author: bdemirb

backends/arm/quantizer/arm_quantizer_utils.py

@@ -476,6 +476,13 @@ class SharedQspecQuantizer(Quantizer, QuantizerReporterUser):
         torch.ops.higher_order.while_loop,
         torch.ops.higher_order.cond,
     ]
+    _UINT8_IO_BRIDGE_OPS: set[Callable[..., object]] = {
+        torch.ops.aten.cat.default,
+        torch.ops.aten.concatenate.default,
+        torch.ops.aten.stack.default,
+        torch.ops.aten.pixel_shuffle.default,
+        torch.ops.aten.slice.Tensor,
+    }
 
     def __init__(self, targets: Optional[list[Callable[..., object]]] = None) -> None:
         super().__init__()
@@ -565,6 +572,32 @@ def _is_quantized_io_boundary(self, node: Node) -> bool:
         """
         return node.op in ("placeholder", "output") and self._is_annotated(node)
 
+    def _qspec_contains_uint8(self, qspec: Any) -> bool:
+        if isinstance(qspec, list):
+            return any(self._qspec_contains_uint8(element) for element in qspec)
+        return getattr(qspec, "dtype", None) == torch.uint8
+
+    def _is_uint8_quantized_io_boundary(self, node: Node) -> bool:
+        if node.op not in ("placeholder", "output") or not self._is_annotated(node):
+            return False
+
+        annotation = node.meta.get(Q_ANNOTATION_KEY)
+        if annotation is None:
+            return False
+
+        if self._qspec_contains_uint8(annotation.output_qspec):
+            return True
+
+        return any(
+            self._qspec_contains_uint8(qspec)
+            for qspec in annotation.input_qspec_map.values()
+        )
+
+    def _model_has_uint8_io(self, model: torch.fx.GraphModule) -> bool:
+        return any(
+            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]:
         shared_nodes = set()
         bfs_queue = [root_node]
@@ -701,9 +734,20 @@ def _annotate_shared_cluster(self, root_node: Node) -> None:
         return
 
     def annotate(self, model: torch.fx.GraphModule) -> None:  # type: ignore[override]
-        for node in model.graph.nodes:
-            if node.target in self.targets and not self._is_annotated(node):
-                self._annotate_shared_cluster(node)
+        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
 
     def validate(self, model: torch.fx.GraphModule) -> bool:  # type: ignore[override]
         return True

backends/arm/test/quantizer/test_uint8_io_quantization.py

@@ -42,6 +42,20 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return torch.clone(x)
 
 
+class CatWithHighRangeBranch(torch.nn.Module):
+    def forward(self, img0: torch.Tensor, img1: torch.Tensor) -> torch.Tensor:
+        image = torch.cat([img0, img1], dim=1)
+        high_range = image * 20.0
+        merged = torch.cat([image, high_range], dim=1)
+        return torch.clone(merged)
+
+
+def _get_observer_scale(prepared, observer_node_name: str) -> float:
+    observer = prepared.get_submodule(observer_node_name)
+    scale, _ = observer.calculate_qparams()
+    return float(scale)
+
+
 def test_uint8_io_quantization_config_tosa_INT_applies_to_io():
     model = SimpleMLP().eval()
     test_data = (torch.rand(1, 4),)
@@ -94,3 +108,29 @@ def test_io_boundary_shared_cluster_is_quantized():
     assert (
         clone_node.meta[Q_ANNOTATION_KEY].output_qspec is not None
     ), "clone node has no output_qspec — IO-boundary cluster stayed in float"
+
+
+def test_cat_does_not_bridge_shared_qspec_clusters():
+    """Regression: cat must not merge image IO and high-range activations into
+    one fallback shared-qspec observer clique.
+    """
+    model = CatWithHighRangeBranch().eval()
+    test_data = (torch.rand(1, 3, 8, 8), 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)
+    prepared(*test_data)
+
+    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))
+
+    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