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