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.

⚠️ 12 Awaiting Approval

As of commit 50754fc with merge base 70594d3 ():

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

• Build documentation (gh)
• Build Presets (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'