Bound the dynamic-qdq traceback in XNNPACK ChannelsLastTaggedReshapePass

[Hyungkeun-Park-Nota]

Summary

The dynamic-quant branch of ChannelsLastTaggedReshapePass.input_to_nhwc traces back over the
q/dq wrapper so the NHWC copy is inserted ahead of the quantize, keeping the x -> q -> dq -> conv
chain XNNPACK matches intact. The loop stops on "args[0] is not a Node" rather than on "this is
not a q/dq node", so it does not stop at the quantized tensor and continues into ordinary compute:

while getattr(input_node, "args", None) and isinstance(
    input_node.args[0], torch.fx.Node
):
    input_node = input_node.args[0]

The input_node.replace_all_uses_with(input_node_nhwc) that follows is then applied from wherever
the walk landed, rewriting consumers the pass never reasoned about. That shows up two ways:

If the walk stops on an intermediate op, lowering succeeds but the serialized graph is
inconsistent, and only the first execute() reports it:

[XNNExecutor.cpp:133] Internal Error: Propagating input shapes failed with code: xnn_status_invalid_parameter
[method.cpp:1421] CALL_DELEGATE execute failed at instruction 3: 0x1

The failing partition contains an elementwise op whose input and output dims disagree, next to a
correct one in the sibling branch:

[0] XNNStaticTranspose  (0)[1,32,160,160] -> (1)[1,160,160,32]
[1] XNNSigmoid          (1)[1,160,160,32] -> (2)[1,32,160,160]
[2] XNNStaticTranspose  (2)[1,32,160,160] -> (3)[1,160,160,32]
[3] XNNSigmoid          (5)[1,32,160,160] -> (4)[1,32,160,160]

If the walk reaches a non-4D constant it fails earlier, during the pass:

RuntimeError: required rank 4 tensor to use channels_last format
While executing _to_copy(%b_mul_10_const_input, memory_format=channels_last)

b_mul_10_const_input is a (256, 1, 1) per-channel constant. A placeholder has no args, so the
walk stops there and tries to convert it. can_be_converted_to_nhwc does check rank 4, but this
path never calls it.

Both were found on w8a8 dynamic-quantized vision models, a YOLOX detector for the first and a SAM
image encoder for the second.

Fix

Restrict the walk to q/dq nodes. dq -> q -> source is two hops and the source is not a q/dq node,
so the walk stops at the source, which is the node it was aiming for, and the
x -> _to_copy -> q -> dq -> conv ordering is unchanged. is_quant is already exported from
backends/xnnpack/utils/quant_utils.py alongside the is_dynamic_qdq this file imports.

Instrumenting the loop on the YOLOX graph: 83 invocations, every one with exactly 2 q/dq hops, and
69 of them (83%) walking past that, up to 26 hops. Bounding the walk leaves the delegate count at
16 either way and removes 16 XNNStaticTranspose nodes (414 to 398 total), so the overshoot was
not buying larger fused partitions.

Test

test_dq_conv2d_eltwise_source_channels_last_tagged_reshape_pass builds the smallest graph that
triggers it: a dynamically quantized conv whose input is a sigmoid reading the placeholder, so
the conv sees sigmoid -> q -> dq. It asserts the sigmoid keeps its placeholder input and that the
channels-last copy sits on the sigmoid's output.

Without the fix the pass produces x -> _to_copy(channels_last) -> sigmoid -> q -> dq -> conv and
the test fails with AssertionError: 'call_function' != 'placeholder'; with it the order is
x -> sigmoid -> _to_copy(channels_last) -> q -> dq -> conv. The assertion is structural because
channels_last does not change eager results, so run_method_and_compare_outputs alone cannot
catch this.

pytest backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py -k eltwise_source

The full file passes (22 tests). Separately, 10 w8a8 dynamic models that already lowered and
executed correctly before this change (googlenet, inception_v3, efficientnet_b4, wideresnet50,
sesr_m5, mobile_vit_s, swin_t, vit_b_16, quicksrnet_small, squeezenet1_0) were re-lowered with the
grouped partitioner and executed: 10/10 pass. The two models above now lower and execute, with
output shapes matching a per-op-partitioned reference build.

cc @GregoryComer @digantdesai @cbilgin @JakeStevens

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/21637

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

⚠️ 13 Awaiting Approval

As of commit 50e14c9 with merge base 70594d3 ():

AWAITING APPROVAL - The following workflows need approval before CI can run:

• Build documentation (gh)
• Build Presets (gh)
• Cadence Build & Test (gh)
• Lint (gh)
• periodic (gh)
• pull (gh)
• Test ARM Backend (gh)
• Test CoreML Backend (gh)
• Test QNN Backend (gh)
• Test QNN Windows MSVC build (gh)
• Test Vulkan Backend (gh)
• Test WebGPU Backend (gh)
• Test XNNPACK Backend (gh)

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[Hyungkeun-Park-Nota]

@pytorchbot label 'module: xnnpack' 'release notes: xnnpack'

[JakeStevens]

Can you add a test for a producer with multiple consumers?

replace_all_uses_with on that shared producer can result in the same class of bug

[JakeStevens]

Should this also be is_dynamic_qdq instead?

[Hyungkeun-Park-Nota]

@JakeStevens Switched the walk guard to is_dynamic_qdq as suggested, and added a multi-consumer test.

That test turned up one more thing: the graph output was being redirected to the channels-last copy, so a fix for that is in here too.

[JakeStevens]

The helper can still rewrite every ordinary sibling consumer. For:

act = torch.sigmoid(x)
sibling = torch.tanh(act)
return conv(act), sibling

the pass rewrites tanh to consume the NHWC copy after tanh was classified as NCHW

[Hyungkeun-Park-Nota]

Thanks. To pass that case as well, changed the helper to redirect only the quantize wrapper (and the choose_qparams feeding it) to the NHWC copy. So, the other consumers keep reading the source. Added the sigmoid/tanh case as a test. Please take a look.

[JakeStevens]

I think may need to switch the approach to instead tracking the subgraph chain we are interested in. as is, this still does not solve the sibling problem fully, as there is an edge case with two siblings that can both be dynamically quantized, eg

class SharedLinearConvDynamicQuant(torch.nn.Module):
      def __init__(self):
          super().__init__()
          self.linear = torch.nn.Linear(8, 8)
          self.conv = torch.nn.Conv2d(3, 4, 1)

      def forward(self, x):
          act = torch.sigmoid(x)
          # Keep linear before conv so it is processed first.
          linear = self.linear(act)

And this cannot really be disambiguated.

So, instead track the specific q/dq chain being traversed, then redirect only its quantize node and associated qparam nodes

[JakeStevens]

thank you for this PR and the back and forth to clean it up!

[Hyungkeun-Park-Nota]

Hi @JakeStevens. Just checking if there’s anything else needed for this PR to be merged. Thanks.