Repository navigation
Check dim order in the optimized layer_norm as the portable one does - #21866
Merged
JakeStevens merged 3 commits intoSep 10, 2026
Merged
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/21866
Note: Links to docs will display an error until the docs builds have been completed.
|
Contributor
Author
|
@pytorchbot label "release notes: ops & kernels" |
SuryanshSS1011
force-pushed
the
fix/optimized-layer-norm-dim-order
branch
from
August 15, 2026 03:48
05bbe79 to
128d5cb
Compare
JakeStevens
reviewed
Aug 18, 2026
| tensors_have_same_dim_order(input, out, mean_out, rstd_out), | ||
| InvalidArgument, | ||
| ret_val); | ||
|
|
Contributor
There was a problem hiding this comment.
Copy the portable op's check for weight and bias too
Contributor
Author
There was a problem hiding this comment.
Done, added the weight and bias checks, plus a test for each so they are covered independently.
Contributor
|
Hi @JakeStevens, is this good to merge? |
JakeStevens
pushed a commit
that referenced
this pull request
Sep 10, 2026
…22355) ### Summary Follow-up to #21866, which did the same for the optimized `layer_norm`. `opt_mm_out` and `opt_bmm_out` carry no dim order check, while portable `mm_out` and `bmm_out` each carry two: ```cpp ET_KERNEL_CHECK( ctx, tensors_have_same_dim_order(in, mat2, out), InvalidArgument, out); ET_KERNEL_CHECK(ctx, tensor_is_default_dim_order(in), InvalidArgument, out); ``` The shared `check_mm_args` and `check_bmm_args` do not supply them either, so choosing the optimized kernels quietly drops a check the portable path enforces. `opt_mm_out` passes raw data pointers to `cpublas::gemm` with the leading dimensions taken from sizes rather than strides, which describes a row-major matrix only in the default dim order. `opt_bmm_out` does the same per batch. Thus, the arithmetic also needs it. This PR adds portable guards to both the kernels, placed in the same position relative to the resize. Not reachable from export today, in the same way the `select_scatter` half of #21915 was not, since `_get_channels_last_dim_order` raises for any rank other than 4 and so a 2D `mm` or 3D `bmm` input cannot be given a non-default dim order. It is still worth adding because the optimized kernels are currently less strict than the portable ones with nothing recording that as intentional, and because #16429 asks for these layouts to be supported, at which point the optimized path would return wrong numbers where the portable path errors. ### Test plan A `NonDefaultDimOrderDies` per op. `op_mm_test.cpp` and `op_bmm_test.cpp` are both in `_optimized_kernels_test_sources`, so the same test body runs against both kernel sets, which makes the gap visible directly: | target | optimized kernels | mm and bmm | |---|---|---| | `optimized_kernels_test` | guards reverted, tests kept | both fail | | `optimized_kernels_test` | with the guards | both pass | | `portable_kernels_test` | untouched either way | both pass | All three tensors in each test share the same non-default dim order, so the same dim order check passes and only the default dim order check can reject. Both skip under `is_aten`. cc @larryliu0820 @manuelcandales @JakeStevens
JakeStevens
pushed a commit
that referenced
this pull request
Sep 10, 2026
…ck (#21865) ### 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: ```cpp 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: ```cpp 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. `slice_scatter` needed `src` in its same dim order check as well, since the kernel reads `src` as one contiguous run and only `input` and `out` were covered. `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. `native_batch_norm` uses the same helpers but is not affected, since it checks `is_contiguous_dim_order` on the input directly. `select_scatter` and `unfold_copy` are affected and are left to a follow-up. `select_scatter` carries only the same dim order check, which passes when a rank-5 input and a rank-4 `src` are both channels-last, because `is_channels_last_dim_order` accepts 4 and 5 dims. `unfold_copy` has no dim order check at all. 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`. It is registered in `targets.bzl` but missing from the CMake list, so those tests do not run in a CMake build and the new one would not either. Ten other op tests are missing the same way, and I have left those alone to keep this to one change. cc @larryliu0820 @manuelcandales @JakeStevens
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Related to #21828 and #21865 (same underlying assumption in the optimized kernel library).
aten.native_layer_normreturns 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:
The optimized kernel has neither that check nor the matching
tensors_have_same_dim_order, andnative_layer_norm.outis mapped totorch::executor::opt_native_layer_norm_outinoptimized.yaml. So any build withEXECUTORCH_BUILD_KERNELS_OPTIMIZED=ONgets the ungated one.It needs the guard rather than stride-aware indexing, because it splits the buffer into
Mrows ofNcontiguous elements: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.NonDefaultDimOrderDiespasses a channels-last input without,meanandrstdall channels-last, so the same dim order check passes and only the default dim order check can reject.meanandrstdshare 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.
cc @larryliu0820 @manuelcandales @JakeStevens