Skip to content

Index portable kernels by the tensor's strides instead of assuming row-major - #21828

Open
SuryanshSS1011 wants to merge 7 commits into
pytorch:mainfrom
SuryanshSS1011:fix/dim-order-gate-flip-gather-scatter
Open

SuryanshSS1011 wants to merge 7 commits into
pytorch:mainfrom
SuryanshSS1011:fix/dim-order-gate-flip-gather-scatter

Conversation

@SuryanshSS1011

@SuryanshSS1011 SuryanshSS1011 commented Aug 14, 2026 •

Copy link
Copy Markdown
Contributor

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:

kernel contiguous input channels-last input
aten.flip matches wrong, max abs diff 32
aten.gather matches wrong, max abs diff 30
aten.scatter matches wrong, max abs diff 30
aten.scatter_add matches rejected, status 0x12
aten.roll matches rejected, status 0x12
aten.permute_copy matches rejected, status 0x12

The cause is one line in coordinateToIndex (runtime/core/exec_aten/util/tensor_util.h):

index += coordinate[d] * getTrailingDims(tensor, d);

getTrailingDims is the product of the sizes after d, 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, and memoizeTrailingDims likewise 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, since getTrailingDims recomputes a product with an overflow check per dimension.

That alone fixes gather, scatter and scatter_add. flip, roll and permute_copy also wrote their output linearly by flat index, which is only correct when the output is contiguous, so those writes are mapped too. The guard on scatter_add's index existed only because this indexing was row-major, so it goes.

getTrailingDims itself is unchanged. Elsewhere it means the length of a contiguous run, for the block copies in cat, split and others, which share the defect but need a different fix.

I took this over adding a tensor_is_default_dim_order guard 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.

cc @larryliu0820 @manuelcandales @JakeStevens

Copilot AI lite review requested due to automatic review settings August 14, 2026 01:51
@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Aug 14, 2026
@pytorch-bot

pytorch-bot Bot commented Aug 14, 2026 •

Copy link
Copy Markdown

🔗 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.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

❌ 2 Cancelled Jobs, 1 Unclassified Failure

As of commit b5ec7d0 with merge base 79e3eae (image):

UNCLASSIFIED FAILURE - DrCI could not classify the following job because the workflow did not run on the merge base. The failure may be pre-existing on trunk or introduced by this PR:

  • Cadence Build & Test / Resolve CI docker image / resolve (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
    ##[error]Refusing to check out fork pull request code from a 'pull_request_target' workflow. This workflow runs with the base repository's GITHUB_TOKEN, secrets, default-branch cache scope, and runner access. Fetching and executing a fork's code in that trusted context commonly leads to "pwn request" vulnerabilities. To opt in, review the risks at https://gh.io/securely-using-pull_request_target and set 'allow-unsafe-pr-checkout: true' on the actions/checkout step.

CANCELLED JOBS - The following jobs were cancelled. Please retry:

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@SuryanshSS1011

Copy link
Copy Markdown
Contributor Author

@pytorchbot label "release notes: ops & kernels"

@pytorch-bot pytorch-bot Bot added the release notes: ops & kernels Changes to the opset and any new / changed kernel implementations label Aug 14, 2026

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 to flip_out, gather_out, and scatter_{src,value}_out portable CPU kernels.
  • Added regression tests asserting kernel failure on channels-last-like inputs for flip, gather, and scatter.
  • Updated op_flip_test.cpp to use an OperatorTest fixture where kernel-failure assertions require access to context_.

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_out still relies on indexToCoordinate(index, ...) and coordinateToIndex(out, ...) (default-layout indexing). Only checking tensor_is_default_dim_order(in) can still allow corruption when index or out are non-default. Add checks for index default dim order and require in/out share dim order (so in default implies out default).
  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.

Comment thread kernels/portable/cpu/op_gather.cpp Resolved
Comment thread kernels/portable/cpu/op_scatter.cpp Resolved
Comment thread kernels/test/op_scatter_test.cpp Resolved
Copilot AI review requested due to automatic review settings August 14, 2026 02:20

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

  • NonDefaultDimOrderDies currently uses a default-dim-order index, so it will fail at the new tensors_have_same_dim_order(in, index, out) check even without the tensor_is_default_dim_order(in) gate. Make index channels-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

  • NonDefaultDimOrderDies currently fails because index is default dim order while self/out are channels-last, so it exercises the new tensors_have_same_dim_order(...) check rather than the intended tensor_is_default_dim_order(in) gate. Make index channels-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

  • NonDefaultDimOrderDies currently fails due to mixed dim orders (index/src are default while self/out are channels-last), so it doesn't validate the new tensor_is_default_dim_order(in) guard. Use channels-last index/src so 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 for roll_out / permute_copy_out), but the PR title/summary/test plan only call out flip, gather, and scatter. 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);

Copilot AI review requested due to automatic review settings August 14, 2026 02:32

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

@SuryanshSS1011 SuryanshSS1011 changed the title Reject non-default dim order in flip, gather, and scatter Reject non-default dim order in the portable kernels that use row-major index math Aug 14, 2026
@SuryanshSS1011
SuryanshSS1011 marked this pull request as draft August 14, 2026 04:22
@SuryanshSS1011
SuryanshSS1011 force-pushed the fix/dim-order-gate-flip-gather-scatter branch from af91b42 to 035e6b6 Compare August 14, 2026 04:39
@SuryanshSS1011 SuryanshSS1011 changed the title Reject non-default dim order in the portable kernels that use row-major index math Index portable kernels by the tensor's strides instead of assuming row-major Aug 14, 2026
@SuryanshSS1011
SuryanshSS1011 marked this pull request as ready for review August 14, 2026 13:23
Copilot AI review requested due to automatic review settings August 14, 2026 13:23

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copilot AI review requested due to automatic review settings August 14, 2026 13:23

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copilot AI review requested due to automatic review settings August 15, 2026 03:39

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copilot AI review requested due to automatic review settings August 18, 2026 15:16

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

@JakeStevens

Copy link
Copy Markdown
Contributor

Seems these changes impact some cadence kernels that use the helper, cc @mcremon-meta

for (size_t i = 0; i < out.numel(); ++i) {
const size_t in_index =
executorch::runtime::coordinateToIndexWithTrailingDimsMemo(
in, in_coord, trailing_dims_memo);
std::memcpy(
out_data + i * element_size,
in_data + in_index * element_size,
element_size);

@mcremon-meta

Copy link
Copy Markdown
Contributor

Seems these changes impact some cadence kernels that use the helper, cc @mcremon-meta

for (size_t i = 0; i < out.numel(); ++i) {
const size_t in_index =
executorch::runtime::coordinateToIndexWithTrailingDimsMemo(
in, in_coord, trailing_dims_memo);
std::memcpy(
out_data + i * element_size,
in_data + in_index * element_size,
element_size);

thanks for the heads up, we don't use dim order so everything should always be contiguous and we should be fine. If CI doesn't break (including the internal Meta CI, which could trigger a revert of this), no objection from my side.

@nil-is-all nil-is-all added the module: kernels Issues related to kernel libraries and utilities, and code under kernels/ label Aug 18, 2026
JakeStevens pushed a commit that referenced this pull request Sep 10, 2026
…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.
@nil-is-all

Copy link
Copy Markdown
Contributor

Seems these changes impact some cadence kernels that use the helper, cc @mcremon-meta

for (size_t i = 0; i < out.numel(); ++i) {
const size_t in_index =
executorch::runtime::coordinateToIndexWithTrailingDimsMemo(
in, in_coord, trailing_dims_memo);
std::memcpy(
out_data + i * element_size,
in_data + in_index * element_size,
element_size);

thanks for the heads up, we don't use dim order so everything should always be contiguous and we should be fine. If CI doesn't break (including the internal Meta CI, which could trigger a revert of this), no objection from my side.

@mcremon-meta could you check again now that CI looks good?

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
@executorch-triage executorch-triage Bot added the community: contribution PRs coming from community (excluding hardware partners) label Sep 22, 2026

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. community: contribution PRs coming from community (excluding hardware partners) module: kernels Issues related to kernel libraries and utilities, and code under kernels/ release notes: ops & kernels Changes to the opset and any new / changed kernel implementations

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants