diff --git a/backends/arm/quantizer/arm_quantizer_utils.py b/backends/arm/quantizer/arm_quantizer_utils.py index a59ccff87b1..b15f01d11c8 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,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 diff --git a/backends/arm/test/quantizer/test_uint8_io_quantization.py b/backends/arm/test/quantizer/test_uint8_io_quantization.py index 3b839dc01c0..0344799321f 100644 --- a/backends/arm/test/quantizer/test_uint8_io_quantization.py +++ b/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