From e52c45eaeb423d3dcad22a9262e1805a546e9981 Mon Sep 17 00:00:00 2001 From: Baris Demir Date: Thu, 25 Jun 2026 15:49:50 +0100 Subject: [PATCH 1/2] Arm backend: Avoid shared qspec bridges for uint8 IO 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 quantization 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. Signed-off-by: Baris Demir Change-Id: Ieeb1cf59e606674ac34050ebe21fc080740cc567 --- backends/arm/quantizer/arm_quantizer_utils.py | 50 +++++++++++++++++-- .../quantizer/test_uint8_io_quantization.py | 40 +++++++++++++++ 2 files changed, 87 insertions(+), 3 deletions(-) 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 From 9245362941229408fac55bfa4cf82d7c491590eb Mon Sep 17 00:00:00 2001 From: Baris Demir Date: Thu, 25 Jun 2026 22:01:14 +0100 Subject: [PATCH 2/2] Arm backend: Keep internal shared qspec with uint8 IO The uint8 IO shared-qspec guard previously removed bridge ops from SharedQspecQuantizer for the whole graph when a model had preserved uint8 IO. That protects image IO accuracy, but it also prevents internal cat clusters from sharing qparams. In RIFE this creates full-resolution FP32/QDQ concat islands before quantized convs. Keep the uint8 IO guard local to shared clusters that actually touch uint8 model IO. Internal shared-qspec clusters can still be annotated and lowered as quantized ops. Add regression coverage for both cases: cat must not bridge uint8 IO to high-range activations, while internal cat between quantized convs must still share qparams. Signed-off-by: Baris Demir Change-Id: I528572d6136f78cbfe27c75579c8b37a4e560887 --- backends/arm/quantizer/arm_quantizer_utils.py | 53 ++++++++++++------- .../quantizer/test_uint8_io_quantization.py | 45 +++++++++++++++- 2 files changed, 77 insertions(+), 21 deletions(-) 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