Arm backend: Keep internal shared qspec with uint8 IO

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,11 +572,40 @@ def _is_quantized_io_boundary(self, node: Node) -> bool:
         """
         return node.op in ("placeholder", "output") and self._is_annotated(node)
 
-    def _get_shared_clique(self, root_node: Node) -> tuple[set[Node], list[Any], bool]:
+    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, 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)
@@ -579,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(
@@ -640,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
@@ -662,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
 

backends/arm/test/quantizer/test_uint8_io_quantization.py

@@ -42,6 +42,33 @@ 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)
+
+
+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()
+    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 +121,57 @@ 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))
+    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 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