Index portable kernels by the tensor's strides instead of assuming row-major - #21828
Index portable kernels by the tensor's strides instead of assuming row-major#21828SuryanshSS1011 wants to merge 6 commits into
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/21828
Note: Links to docs will display an error until the docs builds have been completed.
|
|
@pytorchbot label "release notes: ops & kernels" |
There was a problem hiding this comment.
Pull request overview
This PR addresses silent data corruption in the portable CPU implementations of flip, gather, and scatter when tensors use non-default (e.g., channels-last) dimension order. The approach is to add a default-dim-order guard so these kernels fail fast with InvalidArgument instead of producing incorrect results, and to add regression tests covering the new failure mode.
Changes:
- Added
tensor_is_default_dim_order(...)gating toflip_out,gather_out, andscatter_{src,value}_outportable CPU kernels. - Added regression tests asserting kernel failure on channels-last-like inputs for
flip,gather, andscatter. - Updated
op_flip_test.cppto use anOperatorTestfixture where kernel-failure assertions require access tocontext_.
Reviewed changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| kernels/portable/cpu/op_flip.cpp | Adds a default-dim-order guard to prevent incorrect indexing on channels-last tensors. |
| kernels/portable/cpu/op_gather.cpp | Adds a default-dim-order guard for gather_out. |
| kernels/portable/cpu/op_scatter.cpp | Adds a default-dim-order guard for scatter_src_out and scatter_value_out. |
| kernels/test/op_flip_test.cpp | Adds a regression test expecting failure on non-default dim order using an OperatorTest fixture. |
| kernels/test/op_gather_test.cpp | Adds a regression test expecting failure on non-default dim order. |
| kernels/test/op_scatter_test.cpp | Adds a regression test expecting failure on non-default dim order for scatter_src_out. |
Suppressed comments (1)
kernels/portable/cpu/op_scatter.cpp:148
scatter_value_outstill relies onindexToCoordinate(index, ...)andcoordinateToIndex(out, ...)(default-layout indexing). Only checkingtensor_is_default_dim_order(in)can still allow corruption whenindexoroutare non-default. Add checks forindexdefault dim order and requirein/outshare dim order (soindefault impliesoutdefault).
ET_KERNEL_CHECK(ctx, tensor_is_default_dim_order(in), InvalidArgument, out);
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 12 out of 12 changed files in this pull request and generated no new comments.
Suppressed comments (4)
kernels/test/op_scatter_test.cpp:693
NonDefaultDimOrderDiescurrently uses a default-dim-orderindex, so it will fail at the newtensors_have_same_dim_order(in, index, out)check even without thetensor_is_default_dim_order(in)gate. Makeindexchannels-last too so the test asserts rejection of a channels-last input even when all tensors share the same non-default dim order.
Tensor index = tf_index.zeros({1, 1, 2, 2});
kernels/test/op_gather_test.cpp:387
NonDefaultDimOrderDiescurrently fails becauseindexis default dim order whileself/outare channels-last, so it exercises the newtensors_have_same_dim_order(...)check rather than the intendedtensor_is_default_dim_order(in)gate. Makeindexchannels-last too so all tensors share the same non-default dim order and the test asserts the default-dim-order rejection path.
Tensor index = tf_index.zeros({1, 1, 2, 2});
kernels/test/op_scatter_test.cpp:680
NonDefaultDimOrderDiescurrently fails due to mixed dim orders (index/srcare default whileself/outare channels-last), so it doesn't validate the newtensor_is_default_dim_order(in)guard. Use channels-lastindex/srcso all tensors share the same non-default dim order and the test specifically covers the default-dim-order rejection.
This issue also appears on line 693 of the same file.
Tensor index = tf_index.zeros({1, 1, 2, 2});
Tensor src = tf_data.ones({1, 1, 2, 2});
kernels/portable/cpu/op_scatter_add.cpp:74
- This PR also adds a default-dim-order rejection to
scatter_add_out(and adds new non-default-dim-order tests forroll_out/permute_copy_out), but the PR title/summary/test plan only call outflip,gather, andscatter. Consider updating the PR metadata to reflect the expanded scope so downstream users aren't surprised by additional operators rejecting channels-last inputs.
ET_KERNEL_CHECK(
ctx, tensors_have_same_dim_order(self, src, out), InvalidArgument, out);
ET_KERNEL_CHECK(
ctx, tensor_is_default_dim_order(index), InvalidArgument, out);
ET_KERNEL_CHECK(ctx, tensor_is_default_dim_order(self), InvalidArgument, out);
af91b42 to
035e6b6
Compare
Summary
Three portable kernels accept a channels-last input, return a tensor correctly labelled channels-last, and fill it with wrong data. Nothing errors. Three more share the defect but are currently masked by an unrelated check. Verified all against eager PyTorch with the same op and values and only the memory format differing:
aten.flipaten.gatheraten.scatteraten.scatter_addaten.rollaten.permute_copyThe cause is one line in
coordinateToIndex(runtime/core/exec_aten/util/tensor_util.h):index += coordinate[d] * getTrailingDims(tensor, d);getTrailingDimsis the product of the sizes afterd, which is the row-major stride. A(2, 3, 4, 4)channels-last tensor has strides[48, 1, 12, 3], but that gives[48, 16, 4, 1], so the buffer is walked in the wrong order.This reads
tensor.strides()instead, andmemoizeTrailingDimslikewise for its one caller,permute_copy. For a contiguous tensor the two are equal by definition, so that path is unchanged. It is also cheaper, sincegetTrailingDimsrecomputes a product with an overflow check per dimension.That alone fixes
gather,scatterandscatter_add.flip,rollandpermute_copyalso wrote their output linearly by flat index, which is only correct when the output is contiguous, so those writes are mapped too. The guard onscatter_add'sindexexisted only because this indexing was row-major, so it goes.getTrailingDimsitself is unchanged. Elsewhere it means the length of a contiguous run, for the block copies incat,splitand others, which share the defect but need a different fix.I took this over adding a
tensor_is_default_dim_orderguard to each kernel. A guard would stop the corruption, but the same assumption appears in a dozen more places, and #16429 asks for non-contiguous dim order to be supported rather than rejected. Happy to switch if that's preferred.Test plan
This PR adds seven tests, one per kernel. Each runs a channels-last input and checks it against the result the same values produce contiguously, so it passes only if the kernel honors the tensor's dim order. All seven fail before this change.