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__()