Bound the dynamic-qdq traceback in XNNPACK ChannelsLastTaggedReshapePass

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]

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