Check dim order in the optimized layer_norm as the portable one does

[SuryanshSS1011]

Summary

Related to #21828 and #21865 (same underlying assumption in the optimized kernel library).

aten.native_layer_norm returns wrong data on a channels-last input, off by 3.192 against eager PyTorch where it matches exactly on contiguous input. Nothing errors.

The portable kernel is fine. It carries this, with the reason written down:

// Only support default dim order for now.
// TODO: Support other dim orders.
ET_KERNEL_CHECK(
    ctx, tensor_is_default_dim_order(input), InvalidArgument, ret_val);

The optimized kernel has neither that check nor the matching tensors_have_same_dim_order, and native_layer_norm.out is mapped to torch::executor::opt_native_layer_norm_out in optimized.yaml. So any build with EXECUTORCH_BUILD_KERNELS_OPTIMIZED=ON gets the ungated one.

It needs the guard rather than stride-aware indexing, because it splits the buffer into M rows of N contiguous elements:

const size_t M = getLeadingDims(input, dim);
const size_t N = getTrailingDims(input, dim) * dim_size;

That layout only exists in the default dim order. This copies the two checks across from the portable kernel so the two agree.

Test plan

OpNativeLayerNormTest.NonDefaultDimOrderDies passes a channels-last input with out, mean and rstd all channels-last, so the same dim order check passes and only the default dim order check can reject. mean and rstd share the input's rank with the normalized dims set to 1, which the kernel requires before it reaches any dim order check.

The test file is shared by both kernel libraries, so it runs against the portable kernel too, where it already passes.

[pytorch-bot]

🔗 Helpful Links

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

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

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.

[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"

[Copilot]

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