From 50754fc8ef33682c9435e6b36a9d22348325ea42 Mon Sep 17 00:00:00 2001 From: Hyungkeun Park Date: Fri, 7 Aug 2026 09:47:34 +0000 Subject: [PATCH] Bound the dynamic-qdq traceback in ChannelsLastTaggedReshapePass input_to_nhwc steps back over the dynamic q/dq wrapper so the NHWC copy is inserted ahead of the quantize. The loop stopped only once args[0] was not a Node, so it did not stop at the quantized tensor and ran on into ordinary compute. The rewrite that follows is a blanket replace_all_uses_with from wherever the walk landed, so overshooting either feeds an intermediate op NHWC while leaving that op's own output NCHW, which XNNPACK reports as xnn_status_invalid_parameter when propagating input shapes at execute(), or lands on a non-4D constant and raises "required rank 4 tensor to use channels_last format" in _to_copy. Restrict the walk to q/dq nodes. dq -> q -> source is two hops and the source is not a q/dq node, so it stops there. On a w8a8-dynamic detection model 69 of 83 tracebacks had been overshooting, by up to 26 hops; bounding them leaves the delegate count unchanged and drops 16 now-redundant transposes. --- .../channels_last_tagged_reshape_pass.py | 9 ++-- .../test_channels_last_tagged_reshape.py | 48 +++++++++++++++++++ 2 files changed, 54 insertions(+), 3 deletions(-) diff --git a/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py b/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py index c74c35532b0..8bc5b8ef45c 100644 --- a/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py +++ b/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py @@ -14,6 +14,7 @@ from executorch.backends.xnnpack.utils.quant_utils import ( is_dequant, is_dynamic_qdq, + is_quant, is_tagged_as_implicit_q_dq, tag_as_implicit_q_dq, ) @@ -410,9 +411,11 @@ def input_to_nhwc( is_dynamic_input = is_dynamic_qdq(input_node) if is_dynamic_input: - # Trace back to original source node. Stop if args[0] is not - # a Node (e.g., immutable_list from cat). - while getattr(input_node, "args", None) and isinstance( + # Trace back over the q/dq wrapper to the source node, so the copy + # lands ahead of the quantize. Only q/dq nodes may be stepped over: + # walking further reaches ordinary compute, which the blanket + # replace_all_uses_with below has no business rewriting. + while (is_quant(input_node) or is_dequant(input_node)) and isinstance( input_node.args[0], torch.fx.Node ): input_node = input_node.args[0] diff --git a/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py b/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py index adf1c694b22..758e8109292 100644 --- a/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py +++ b/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py @@ -364,6 +364,54 @@ def test_dq_conv2d_channels_last_tagged_reshape_pass(self) -> None: .run_method_and_compare_outputs() ) + class EltwiseConv2dDynamicQuant(torch.nn.Module): + def __init__(self): + super().__init__() + self.conv = torch.nn.Conv2d(3, 10, 3) + + def forward(self, x): + return self.conv(torch.sigmoid(x)) + + def test_dq_conv2d_eltwise_source_channels_last_tagged_reshape_pass(self) -> None: + # The conv's input is sigmoid -> q -> dq. Stepping past the q/dq pair leaves + # the sigmoid reading NHWC while its own output stays NCHW, which XNNPACK + # only rejects at runtime. + tester = ( + Tester(self.EltwiseConv2dDynamicQuant().eval(), (torch.randn(1, 3, 8, 8),)) + .quantize( + Quantize( + quantization_config=get_symmetric_quantization_config( + is_dynamic=True + ) + ) + ) + .export() + .to_edge() + .run_passes(self.PassStage) + ) + + artifact = tester.get_artifact(StageType.RUN_PASSES) + graph_module = artifact.exported_program().graph_module + sigmoid_nodes = [ + node + for node in graph_module.graph.nodes + if node.target == exir_ops.edge.aten.sigmoid.default + ] + self.assertEqual(len(sigmoid_nodes), 1) + sigmoid = sigmoid_nodes[0] + + # The sigmoid keeps its NCHW input and the copy sits on its output instead. + self.assertEqual(sigmoid.args[0].op, "placeholder") + copies = [ + user + for user in sigmoid.users + if user.target == exir_ops.edge.aten._to_copy.default + and user.kwargs.get("memory_format") == torch.channels_last + ] + self.assertEqual(len(copies), 1) + + tester.run_method_and_compare_outputs() + class ConvAddConvOutput(torch.nn.Module): def __init__(self): super().__init__()