Reject non-default dim order in the portable kernels that copy by block

[SuryanshSS1011]

Summary

Follow-up to #21828, which fixes the coordinate-indexed half of the same defect. These two PRs touch no files in common and can land in either order.

Five portable kernels accept a channels-last input and return wrong data without erroring. Three more share the defect but are currently masked by an unrelated check. Checked against eager PyTorch with only the memory format differing:

kernel	contiguous input	channels-last input
aten.cat	matches	wrong, max abs diff 4.141
aten.pixel_unshuffle	matches	wrong, max abs diff 3.941
aten.slice_scatter	matches	wrong, max abs diff 3.376
aten.topk	matches	wrong, max abs diff 2.581
aten.constant_pad_nd	matches	wrong, max abs diff 2.526
aten.cumsum	matches	rejected, status 0x12
aten.split_copy	matches	rejected, status 0x12
aten.split_with_sizes_copy	matches	rejected, status 0x12

The check that masks the last three is tensors_have_same_dim_order, and it fires only because the exported graph gives them differing dim orders. Nothing guards the arithmetic itself, so the unit tests below reach those kernels directly.

These kernels copy in blocks, taking getLeadingDims and getTrailingDims around the operating dim as the length of a contiguous run:

const size_t outer = getLeadingDims(out, dim);
const size_t dim_stride = getTrailingDims(out, dim);

That run only exists in the default dim order. Under channels-last the elements after dim are not adjacent, so the arithmetic walks the wrong bytes.

This is the same assumption as #21828, but it needs a different fix. There the kernels index element by element, so reading the tensor's strides is enough. Here the algorithm depends on contiguity itself, and supporting the layout would mean restructuring each loop. So this adds the guard that op_addmm, op_bmm and op_avg_pool2d already use:

ET_KERNEL_CHECK(ctx, tensor_is_default_dim_order(in), InvalidArgument, out);

cat needed it on each input as well as the output, since it had no dim order check at all. pixel_shuffle already carries this exact pair of checks, so pixel_unshuffle missing them looks like an oversight rather than a decision.

Note: cat and split appear in nearly every model, so a channels-last program that runs today would start failing instead of returning wrong numbers. That still seems better than silent corruption, but happy to restructure the loops to support the layout instead if that's preferred.

select_scatter and native_batch_norm use the same helpers but are not affected. Their auxiliary tensors are lower rank, so the existing same dim order check already forces the input contiguous. unfold_copy is likely affected, but the wrong result I measured came from a multi-op graph and I could not attribute it to the kernel. The same block arithmetic appears in the optimized layer_norm, fixed in #21866, and in the quantized dequantize, which #21517 covers.

Test plan

A NonDefaultDimOrderDies test per kernel. Each passes a uniformly channels-last set of tensors, so the existing same dim order check passes and only the new guard can reject. Reverting the eight kernel sources while keeping the tests fails all eight.

This also adds op_pixel_unshuffle_test.cpp to kernels/test/CMakeLists.txt. The list includes op_pixel_shuffle_test.cpp and every other op in that directory but omits this one, so those tests do not run in a CMake build and the new one would not either.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 54aa30a with merge base 7dd3335 ():

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.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[SuryanshSS1011]

@pytorchbot label "release notes: ops & kernels"