diff --git a/backends/arm/quantizer/arm_quantizer_utils.py b/backends/arm/quantizer/arm_quantizer_utils.py index a59ccff87b1..a79b2c66e92 100644 --- a/backends/arm/quantizer/arm_quantizer_utils.py +++ b/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 diff --git a/backends/arm/test/quantizer/test_uint8_io_quantization.py b/backends/arm/test/quantizer/test_uint8_io_quantization.py index 3b839dc01c0..95ad768ec66 100644 --- a/backends/arm/test/quantizer/test_uint8_io_quantization.py +++ b/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