Repository navigation
Reject non-default dim order in the portable kernels that copy by block - #21865
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/21865
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 69a6764 with merge base 7dd3335 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
|
@pytorchbot label "release notes: ops & kernels" |
|
@claude please review |
|
For Only input and out are validated. A channels-last src with default-order input/out still passes, then src_data is read linearly as contiguous data, silently permuting values. Change the check to tensors_have_same_dim_order(input, src, out) and add a mixed-layout regression test. |
The kernel reads src as one contiguous run, but only input and out were covered by the same dim order check, so a channels-last src alongside a default-order input and out was accepted and returned permuted values. Fold src into the existing check and add a regression test for it. The uniform channels-last test still fails on the default dim order check rather than this one, so the two tests pin different guards.
|
Thanks for the review and the catch @JakeStevens! I have folded Looking into this further turned up two more kernels with the same defect:
I will work on their fixes via a separate PR. |
|
nit: A few of the tests only exercise the first condition, for example: cat used channel last for input and output in the test the op checks output first, errors out we never check input So we are not really testing that pathway each arg should be tested independently for completeness |
|
Made the change:
|
### Summary Follow-up to #21865. Two more portable kernels index with `getLeadingDims` and `getTrailingDims`, which describe a contiguous run only in the default dim order. `unfold_copy` carries no dim order check at all. Checked against the same logical data in two memory formats: | kernel | contiguous input | channels-last input | |---|---|---| | `aten.unfold_copy` | matches | wrong, max abs diff 14.0 | It reads `self` as a contiguous run and writes `out` as a flat run, so neither side survives a non-default layout. `unfold` appears in shipped example models including the Qualcomm llama vision encoder and the gemma4 speech transform, where a channels-last input is ordinary. `select_scatter` has the same arithmetic and the same defect, but it is not reachable from export today. It carries `tensors_have_same_dim_order(in, src, out)`, which is not sufficient on its own, because `is_channels_last_dim_order` accepts 4 and 5 dims, so a rank-5 input and its rank-4 `src` can both be channels-last and pass together. At rank 4 the `src` is rank 3 and the check does reject. Reaching the rank-5 case needs a channels-last dim order that `exir/dim_order_utils.py` refuses to generate, since it raises for any rank other than 4. So the kernel is protected by a gap in the export layer rather than by its own checks, and the guard closes that before 5D channels-last support makes it reachable. Both get the guard `op_addmm`, `op_bmm` and the kernels in #21865 use: ```cpp ET_KERNEL_CHECK(ctx, tensor_is_default_dim_order(in), InvalidArgument, out); ``` `unfold_copy` needs it on `out` as well, since `out` is one rank above `self` and is written as a flat run. Checking the pair with `tensors_have_same_dim_order` would not work there, for the same rank reason as above. ### Test plan A `NonDefaultDimOrderDies` per kernel, plus a second `unfold_copy` case covering the `out` guard. Reverting the two kernel sources while keeping the tests fails all three. The `select_scatter` test has to be rank 5. A rank-4 case is rejected by the pre-existing same dim order check on current main, so it would pass without reaching the new guard. It builds its tensors with `make_with_dimorder`, because `channels_last_like` asserts `sizes.size() == 4` and cannot express the failing case. All three tests skip under `is_aten`, since the ATen kernels handle a non-default dim order and do not throw. This also registers `op_unfold_copy_test.cpp` in `kernels/test/CMakeLists.txt`, without which neither the existing tests in that file nor the new ones run in a CMake build.
…21866) ### 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: ```cpp // 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: ```cpp 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.
Sorry for the delay, re-running CI |
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:
aten.cataten.pixel_unshuffleaten.slice_scatteraten.topkaten.constant_pad_ndaten.cumsumaten.split_copyaten.split_with_sizes_copyThe 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
getLeadingDimsandgetTrailingDimsaround the operating dim as the length of a contiguous run:That run only exists in the default dim order. Under channels-last the elements after
dimare 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_bmmandop_avg_pool2dalready use:ET_KERNEL_CHECK(ctx, tensor_is_default_dim_order(in), InvalidArgument, out);catneeded it on each input as well as the output, since it had no dim order check at all.slice_scatterneededsrcin its same dim order check as well, since the kernel readssrcas one contiguous run and onlyinputandoutwere covered.pixel_shufflealready carries this exact pair of checks, sopixel_unshufflemissing them looks like an oversight rather than a decision.Note:
catandsplitappear 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_normuses the same helpers but is not affected, since it checksis_contiguous_dim_orderon the input directly.select_scatterandunfold_copyare affected and are left to a follow-up.select_scattercarries only the same dim order check, which passes when a rank-5 input and a rank-4srcare both channels-last, becauseis_channels_last_dim_orderaccepts 4 and 5 dims.unfold_copyhas no dim order check at all. The same block arithmetic appears in the optimizedlayer_norm, fixed in #21866, and in the quantizeddequantize, which #21517 covers.Test plan
A
NonDefaultDimOrderDiestest 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.cpptokernels/test/CMakeLists.txt. It is registered intargets.bzlbut 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