[msluszniak]

Summary

A delegate input whose dim order is neither contiguous nor channels-last (e.g. (1, 0, 2) from a transposed example input) is serialized with its logical sizes, but XNNExecutor reshapes it with sizes in dim order, so execution fails with xnn_status_invalid_parameter or reads swapped axes. Define such inputs in memory order and transpose back at the start of the graph. #14735 was the constant-tensor side of this.

Test plan

New ops/test_io_dim_order.py tests fail before, pass after. backends/xnnpack/test/{ops,passes,serialization} unchanged.

cc @GregoryComer @digantdesai @cbilgin @JakeStevens @nil-is-all

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 1 New Failure

As of commit bf19caf with merge base 4628b2b ():

NEW FAILURE - The following job has failed:

• pull / nxp-mcuxpresso-test / linux-job (gh)
  undefined reference to SYS_INFO_PrintAll'`

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

[JakeStevens]

@claude please review

[JakeStevens]

those CI failures look likely to actually be related

[claude]

Claude finished @JakeStevens's task in 3m 34s —— View job

---

The preprocess fix looks right for float inputs, but the new test feeds scrambled input data, so it fails in CI whether or not the fix works. It needs to be fixed and the coverage widened before this merges.

Detailed review

PR Review: #24081

Summary

The PR declares XNNPACK delegate inputs whose dim order is neither contiguous nor channels-last in memory order, then adds a static transpose at the start of the graph to restore the logical layout. The preprocess logic looks correct: the perm direction is right, the transpose prologue runs before other nodes, and external ids are kept. The main problems:

• The new test fails in CI because the test itself is wrong, so the fix is not yet verified.
• Quantized inputs with these dim orders now make lowering fail outright.
• Outputs have the same problem and are not handled.
• Test coverage is thin.

Code Quality

• backends/xnnpack/xnnpack_preprocess.py:151-156: quantized inputs now abort the whole lowering. The partitioner still claims a quantized placeholder with a dim order like (1, 0, 2). Preprocess then raises RuntimeError, so to_edge_transform_and_lower fails instead of leaving that input outside the delegate. Before this PR the same graph lowered and only went wrong at runtime. Now it is a hard error at export time. Either reject these inputs in the partitioner or handle them; the static transpose works on quantized values too, given matching quant params on the memory-order value. Not exercised by a test.
• Outputs have the same bug and are not handled. XNNExecutor::resize_outputs (runtime/XNNExecutor.cpp:268-270) maps XNNPACK's output dims through the output tensor's dim order, but outputs are still serialized with logical dims. A delegate output whose spec has a non-default, non-NHWC dim order would get permuted sizes. Please add a matching transpose at the end of the graph, or state in the PR that only inputs are fixed. Low confidence: I did not confirm that such an output spec can occur today.
• xnnpack_preprocess.py:95-107: stale assert_default_dim_order. Nothing calls it, and its message ("only supports contiguous memory format for inputs") now contradicts this PR. Delete it.
• xnnpack_preprocess.py:168-170: dead input_ids rewrite. XNNCompiler.cpp:675-676 rebuilds input ids from each value's XNN_VALUE_FLAG_EXTERNAL_INPUT flag, so rewriting xnn_graph.input_ids changes nothing at runtime. Drop it, or keep it only if something else reads it.
• xnnpack_preprocess.py:145: two separate layout mechanisms. NHWC inputs are handled by ChannelsLastTaggedReshapePass on the FX graph. Every other dim order is now patched afterwards on the serialized XNNGraph. The two stay in sync only through the private ChannelsLastTaggedReshapePass._is_nhwc_tensor. Patching the serialized graph is defensible, since an FX placeholder can't carry a memory-order shape. But the coupling should be explicit, for example a shared public helper, so a change to the NHWC rule can't silently send tensors down the wrong path.
• Comment length. The 9-line docstring on define_inputs_in_memory_order and the test docstring repeat the PR description, and the test docstring describes old behavior ("Previously…"). CLAUDE.md asks for minimal comments.

Testing

• Blocker: backends/xnnpack/test/test_xnnpack_partitioner.py:186-190 passes the wrong input.
  • The comment "The bindings copy an input's buffer as is" is false. run_method calls use_layout(input_tensor_meta(i)->dim_order()) (extension/pybindings/pybindings.cpp:2531). For a tensor that isn't already in the exported layout, use_layout copies it into that layout with empty_strided(...).copy_(owner_) (pybindings.cpp:858-870), which keeps the logical values.
  • x.transpose(0, 1).contiguous().view(x.shape) reinterprets the bytes, so its logical values are not x. The delegate therefore computes bmm(x_in_memory_order, y), and the comparison with bmm(x, y) fails. That matches the 4 CI failures.
  • Fix: pass x (or x.contiguous()) directly and delete the comment. Until then, the "fails before, passes after" test plan has not been shown.
• The perm direction isn't tested. (1, 0, 2) is its own inverse, so perm=dim_order and the correct perm=[dim_order.index(i) ...] give identical results. Add a 3-cycle such as (1, 2, 0) or (2, 0, 1) so the test catches a regression in the perm direction.
• Missing cases:
  • static shapes;
  • a dynamic-shape run at a size other than the example size (e.g. n=2 and n=8);
  • 4D non-NHWC orders such as (0, 2, 1, 3);
  • 5D channels_last_3d, which _is_nhwc_tensor rejects, so it now takes the new path;
  • a mix of default and non-default inputs;
  • the quantized-input path.
• Test location and style. test_xnnpack_partitioner.py covers partitioner behavior and deprecation warnings, but this is a preprocess/runtime fix. It fits better in backends/xnnpack/test/ops/ (e.g. test_bmm.py) or the serialization tests. It should use the existing Tester flow ending in .run_method_and_compare_outputs() instead of hand-rolled _load_for_executorch_from_buffer plus allclose.

Recommendation

Request Changes

The new test is wrong and fails on every CI platform, so the fix has no passing coverage yet. Fix the test input and add a non-self-inverse dim order. Also decide what happens to quantized inputs (partitioner rejection rather than a lowering-time RuntimeError) and to outputs.

[msluszniak]

I'm on it, addressing the incorrect test.

[msluszniak]

Ok, the blocker with the incorrect input to the test is fixed together with other things. @JakeStevens, please check again.

[nil-is-all]

Rebase will resolve failing arm CI test. @JakeStevens could you review again?

[msluszniak]

Rebased onto main.

[JakeStevens]

Thanks for the update, the transpose directions and the dynamic-shape handling look right. Two problems with how the dim order is determined, though, and one of them is a regression:

1. val.dim_order() doesn't match the dim order the runtime uses when there are size-1 dims (regression).
XNNExecutor shapes externals by the ETensor's dim order, which the emitter computes with exir.tensor.dim_order_from_stride (a stable sort by stride). torch.Tensor.dim_order() breaks ties on size-1 dims differently. Example: torch.randn(1, 8, 1, 16).transpose(1, 2) has shape (1, 1, 8, 16) and strides (128, 16, 16, 1):

• torch dim_order() gives (0, 2, 1, 3)
• dim_order_from_stride gives (0, 1, 2, 3)
• is_nhwc_tensor is False, because the tensor is both contiguous and channels-last contiguous

So preprocess declares the external as [1, 8, 1, 16] and adds a transpose, while the runtime reshapes it as [1, 1, 8, 16]. On main this input works, because the logical dims happen to match the default order. Batch-1 or seq-len-1 transposes would hit this. Computing the order with dim_order_from_stride(val.stride(), tuple(val.shape)) should fix it.

2. Output dim orders are read after the XNNPACK passes have run.
node_to_external_map is built from the post-XNNPACKPassManager graph. When the producer of an output is NHWC-tagged, ChannelsLastTaggedReshapePass inserts a _to_copy(contiguous_format) before the output, so after the retrace the val is contiguous and no epilogue is emitted. The runtime, however, still uses the original output spec. For example, in forward(x, y): return y + conv(x) with y in dim order (0, 2, 1, 3), eager add follows y, but in the pass add gets tagged NHWC. I traced this through the code but haven't reproduced it. Taking the IO dim orders from edge_program before the passes run, matched by position, would fix both 1 and 2. Also, an external without a val should probably raise instead of being skipped with continue.

Smaller things:

• hasattr(quant_params, "channel_dim"): externals are activations, which are per-tensor quantized, so the per-channel remap looks unreachable and isn't tested. Remapping channel_dim alone also isn't enough for PerChannelGroupQuant. I'd accept only PerTensorQuant and raise for anything else.
• Tests: please add a size-1 dim case and a case where an NHWC op (conv or pool) feeds a non-default output. test_fp32_add_transposed_input_and_output also doesn't assert that the output's dim order is non-default.
• test_qs8_add_quantized_permuted_input compares int8 XNNPACK output against the q/dq reference at atol=1e-5 with unseeded inputs. A single requantization rounding difference is about one output scale, so I'd expect this to flake. It also doesn't check the delegate count or that the delegate input is actually int8.

Review drafted with Claude Code.

[msluszniak]

Dim orders now come from dim_order_from_stride on the program before the XNNPACK passes, matched by position, and an external without a val raises. Only PerTensorQuant is accepted for a transposed quantized external. Added the size-1 and conv-to-transposed-output cases, the output dim order assert, and a seeded qs8 test with a one-step tolerance, delegate count and int8 input checks.

[JakeStevens]

I checked the transpose directions and the runtime reshape and resize paths, and they look right. The one thing I'd like addressed before merging is a transposed input feeding a dynamically quantized op, which now aborts lowering (see the comment on line 154).

Review drafted with Claude Code.

[JakeStevens]

A transposed input that feeds a dynamically quantized op (the 8da4w/8da8w linear path) reaches this raise. For dynamic quant, get_quant_params returns PerTokenDynamicQuant, and XNNCompiler.cpp:410-418 requires that tensor to keep its external id. So the transpose can't be inserted in front of it, and clearing the id at line 163 would break it at runtime anyway.

The partitioner still claims these inputs, so to_edge_transform_and_lower now aborts where it used to lower and only misbehave at runtime. Could the partitioner decline inputs with a non-default dim order that feed a dynamic-quant op? Or could the transpose go before the dynamic-quant convert? If neither fits in this PR, please fix the message ("must be quantized per tensor" is misleading for a float input that is dynamically quantized) and add a test for this path.

[msluszniak]

The external stays fp32: dynamic quant is an internal XNNConvert after the transpose (410-418 is the legacy dq_datatype path, never serialized). Added test_qd8_linear_permuted_input, which fails on main with invalid_parameter.

[JakeStevens]

thanks for taking in all the feedback! ping us if (when..) we forget to check back in on CI