Repository navigation
Make nonzero and the Vulkan, OpenVINO and WebGPU delegates follow the declared dim order - #24771
Open
shoumikhin wants to merge 8 commits into
Open
shoumikhin wants to merge 8 commits into
shoumikhin wants to merge 8 commits into
Conversation
shoumikhin
requested review from
JacobSzwejbka,
larryliu0820 and
mergennachin
as code owners
November 7, 2026 05:38
shoumikhin
added
the
release notes: exir
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/24771
Note: Links to docs will display an error until the docs builds have been completed. ❌ 1 Cancelled Job, 4 Unrelated FailuresAs of commit cf1f16d with merge base e0015e1 ( CANCELLED JOB - The following job was cancelled. Please retry:
BROKEN TRUNK - The following jobs failed but were present on the merge base:👉 Rebase onto the `viable/strict` branch to avoid these failures
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
Contributor
Author
|
@claude /pr-review |
|
Claude finished @shoumikhin's task in 29m 48s —— View job PR Review: #24771Needs discussion: Detailed reviewSummaryThis PR fixes the layout of the
Code Quality
Infrastructure
Testing
API Design
Backward Compatibility
Performance
RecommendationNeeds Discussion The main open question is whether
No code was changed. I did not build or run tests, because running Python here needs approval. Every finding comes from reading the code. All 19 changed files were reviewed. |
shoumikhin
force-pushed
the
nonzero-masked-select-device-io
branch
from
November 7, 2026 18:04
abe9f62 to
ad64232
Compare
shoumikhin
changed the title
Contributor
Author
|
@claude /pr-review |
|
Claude finished @shoumikhin's task in 23m 44s —— View job PR Review: #24771Needs discussion: the Detailed reviewSummaryThis PR makes the portable
The main problem is that the
Backward Compatibility
Code Quality
Testing
RecommendationNeeds Discussion The
I didn't change any code. I didn't build or run tests either, so every finding comes from reading the code. All 27 changed files were reviewed. |
Contributor
Author
|
@claude /pr-review |
shoumikhin
force-pushed
the
nonzero-masked-select-device-io
branch
from
November 7, 2026 19:16
ad64232 to
ee9b608
Compare
|
Claude finished @shoumikhin's task in 33m 34s —— View job PR Review: #24771Needs discussion: the portable Detailed reviewSummaryThe PR makes the portable Code Quality
Testing
Backward Compatibility
Performance
RecommendationNeeds Discussion Before merging, decide which readers that ignore the dim order of the
I didn't change any code. I didn't build or run tests either, so every finding comes from reading the code. The verifier's output-side point is marked low confidence for that reason. All 27 changed files were reviewed. |
shoumikhin
force-pushed
the
nonzero-masked-select-device-io
branch
from
November 8, 2026 00:39
ee9b608 to
1fd85e6
Compare
Contributor
Author
|
Thanks. Changes:
Not changed: a C++ test for the returned external constant (same predicate as the Python test; needs a new generated program), and skipping the copy for 1-D or index-only readers (kernels compare dim orders, not bytes, so |
Contributor
Author
|
@claude /pr-review |
|
Claude finished @shoumikhin's task in 48m 47s —— View job PR Review: #24771Approve with one follow-up: Vulkan, WebGPU and OpenVINO delegates that take a host Detailed reviewSummaryThe PR makes the portable Backward Compatibility
Performance
Code Quality
RecommendationApprove Nothing blocks. The Vulkan, WebGPU and OpenVINO misread comes from a dim-order gap those runtimes already have, which this PR exposes for Progress
Review only: no code was changed, and nothing was built or run. |
A delegate whose graph reads and writes every tensor in the default dim order has to copy an input stored in another dim order, such as channels last or the column-major result of `nonzero`, into a buffer in the default order, and copy an output from such a buffer back into its own dim order. The next changes make the Vulkan, OpenVINO and WebGPU runtimes do this, so this adds the shared code first. `tensor_data_is_in_default_dim_order()` tells whether the elements of a tensor already sit where the default dim order puts them. Like `output_is_raw_copyable` in the CUDA runtime, it ignores the strides of dims of size 1, so a channels-last tensor with a height and a width of 1 needs no copy. `copy_to_default_dim_order()` and `copy_from_default_dim_order()` copy between a tensor and such a buffer. They visit the elements in the default dim order, with a loop over the last dim and an odometer over the other dims that tracks the offset in the data, so they need no division per element. Elements of 1, 2, 4 or 8 bytes are copied with a fixed-size copy. A tensor whose data is already in the default order is copied with one memcpy. The new tests copy a [2, 3, 4] tensor stored in each of its dim orders to and from a buffer, for elements of 1, 2, 4 and 8 bytes and with the element size known only at run time. They also cover scalars, empty tensors and dim orders that differ only in dims of size 1.
The Vulkan graph reads and writes every tensor as if it were stored in the default dim order, but the runtime copied inputs and outputs as flat bytes. So an input in another dim order, such as a column-major `[N, 2]` index, was read in the wrong places, and an output declared in another dim order, such as channels-last, got its data in the default order, with no error. The runtime now copies such an input into the default order before the upload, and such an output from the default order into its own after the download, with `copy_to_default_dim_order()` and `copy_from_default_dim_order()`. A tensor whose data already sits in the default order, such as any tensor in the default dim order, takes the same path as before, with no extra copy. Before, a channels-last input and output came out right only for elementwise ops on tensors of one shape, because both were copied flat. Now the op does not matter. This also applies to method inputs and outputs that go straight to the delegate. A caller that passes a tensor in the default dim order for such an input has its data treated as already in the method's dim order, as portable kernels always did. Such an output now holds its data in its declared dim order. So a caller that relied on the flat copies must now pass and read the data in the method's dim order. The portable `nonzero` kernel still writes its result row by row, although the program declares dim order `(1, 0)` for it. A Vulkan delegate that reads such a result was right only because neither side followed the dim order, so it reads that result wrong until the kernel writes in its declared order. The new tests pass a column-major index to a delegated `gather`, and scale the channels of a channels-last input. Both fail without the fix.
OpenVINO reads and writes every buffer as if it were stored in the default dim order, but the runtime handed it each input and output buffer with only its sizes. So an input in another dim order, such as a column-major `[N, 2]` index, was read in the wrong places, and an output declared in another dim order, such as channels-last, got its data in the default order, with no error. The runtime now copies such an input into a buffer that OpenVINO owns, in the default order, with `copy_to_default_dim_order()`. Such an output goes to a buffer that OpenVINO owns, as an output whose size can change already does, and is copied into its own dim order after inference with `copy_from_default_dim_order()`. A tensor whose data already sits in the default order, such as any tensor in the default dim order, takes the same path as before. Before, a channels-last input and output came out right only for elementwise ops on tensors of one shape, because both were handed over as they were. Now the op does not matter. This also applies to method inputs and outputs that go straight to the delegate. A caller that passes a tensor in the default dim order for such an input has its data treated as already in the method's dim order, as portable kernels always did. Such an output now holds its data in its declared dim order. So a caller that relied on the flat copies must now pass and read the data in the method's dim order. The portable `nonzero` kernel still writes its result row by row, although the program declares dim order `(1, 0)` for it. An OpenVINO delegate that reads such a result was right only because neither side followed the dim order, so it reads that result wrong until the kernel writes in its declared order. The OpenVINO guide and README no longer say that a returned copy made with `.to()`, a copy in another dim order, or a copy of a value computed inside the partition can come back wrong, since each comes back right now. For the same reason, the partitioner now keeps a copy of a constant in the partition also when the copy is in another dim order, as it already did for the other copies of constants, so the program stores the constant once. Its docstring no longer gives the old limits as reasons. The new tests pass a column-major index to a delegated `index_select`, scale the channels of a channels-last input, and return a channels-last copy of a buffer that stays in the partition. They fail without the fix.
The WebGPU graph reads and writes every tensor as if it were stored in the default dim order, but the runtime uploaded inputs and read back outputs as flat bytes. So an input in another dim order, such as a column-major `[N, 2]` index, was read in the wrong places, and an output declared in another dim order, such as channels-last, got its data in the default order, with no error. The runtime now copies such an input into the default order before the upload, with `copy_to_default_dim_order()`. It reads such an output back into a buffer in the default order and copies it into its own dim order with `copy_from_default_dim_order()`. A tensor whose data already sits in the default order, such as any tensor in the default dim order, takes the same path as before, with no extra copy. Before, a channels-last input and output came out right only for elementwise ops on tensors of one shape, because both were copied flat. Now the op does not matter. This also applies to method inputs and outputs that go straight to the delegate. A caller that passes a tensor in the default dim order for such an input has its data treated as already in the method's dim order, as portable kernels always did. Such an output now holds its data in its declared dim order. So a caller that relied on the flat copies must now pass and read the data in the method's dim order. The new native tests feed a column-major index to a delegated `gather`, and scale the channels of a channels-last input. Both fail without the fix. A Python test checks that both models are delegated, so the native tests cannot pass on portable kernels instead.
ATen returns the result of `torch.nonzero` column by column: for a result of shape `[N, ndim]`, the strides are `(1, N)`. Exported programs keep that layout, so the output tensor of `nonzero.out` has dim order `(1, 0)`. The portable kernel ignored this and always wrote row by row, so the runtime read the indices back in the wrong places, with no error: ``` x = [[0, 1, 2], [3, 0, 0]] eager: [[0, 1], [0, 2], [1, 0]] ExecuTorch: [[0, 2], [1, 1], [0, 0]] ``` The kernel now writes each index through the strides of the output, so it is right for both layouts. A resized dynamic output gets its strides recomputed from its dim order, so this also holds for the upper-bounded outputs that exported programs use. The kernel also reads its input through its strides now. A `nonzero` that reads another `nonzero` result stays right this way. Before, it read its input row by row, which matched only the old row-by-row writes. This also fixes the `nonzero` of an input whose data really is in another dim order, such as a channels-last mask. The portable `gather` and `scatter` kernels read their tensors as if they were stored row by row, and `scatter` copied `self` into its output byte by byte. With a `nonzero` result as the index, as the `self` of either op, or as the `src` of `scatter`, they were right only because of the old row-by-row writes. They now read the index, `self` and `src` through their strides, as `index` already does for its indices, and `scatter` copies `self` into its output through the strides of both. A contiguous index is still read in order, as before. `unsqueeze_copy` copied its input into its contiguous output byte by byte, so it had the same problem. It now copies an input in another dim order element by element, in logical order, and refuses an output that is not in the default dim order, as `view_copy` already does. `narrow_copy` read its input row by row too. The slice copy behind it now copies element by element, through the strides of both tensors, unless its block copy is right for them: both tensors have the same dim order, and it starts with dims 0 to dim, in order. That holds for any two contiguous tensors and for two channels-last tensors sliced along dim 0, so those keep the block copy. Two channels-last tensors sliced along another dim used to take the block copy too, so `slice_copy` and `narrow_copy` of a channels-last tensor along C, H or W gave wrong values, and so did a `narrow_copy` into an output in another dim order. They take the strided copy now. The portable and Cadence HiFi `slice_copy` refuse mixed dim orders before they reach that copy, as before. The Cadence Fusion G3 `slice_copy` checks dim orders only when built with OP_ARG_CHECK, so for the dtypes it does not pass to `xa_nn_slice`, such as Long, it now gets the strided copy too, or an error when its output does not have the shape of the slice. So a program where `nonzero` feeds `gather`, `scatter`, `unsqueeze_copy` or `narrow_copy` keeps working, including programs exported before this change and programs exported with `_skip_dim_order=True`. A `nonzero` result now has the `(1, 0)` layout that the program declares. C++, Java or Objective-C code that reads it row by row, and a delegate that takes it as an input, must follow its dim order, as the CUDA, Core ML, Vulkan, WebGPU and OpenVINO backends do. These kernels read any input in its declared dim order now, not only a `nonzero` result. So an input that declares another dim order but holds its data in the default order is read differently than before. A delegate that writes an output in the default order, whatever dim order the output declares, produces such an input. The new kernel tests fail without the fix.
`scatter.src` and `scatter_add` accept a 0-dim index with a 1-dim src, as PyTorch does: ``` torch.scatter(torch.zeros(3), 0, torch.tensor(1), torch.tensor([5., 6., 7.])) # tensor([0., 5., 0.]) ``` The portable kernels find the src element of each index element from the coordinate of that index element. A 0-dim index has no coordinate, so they read the offset in src from uninitialized memory and gave wrong values or crashed, depending on what the stack held. For `scatter.src` this happens with a 0-dim or a 1-dim self, and for `scatter_add` with a 1-dim self. `gather` already handles a 0-dim index on its own, and `scatter.src` and `scatter_add` now do the same: they write, or add, the first element of src at the position the index holds. `scatter.value` does not need this, since it reads no src and its output then has at most one dim. Without this change the new tests read the uninitialized coordinate, so whether they fail depends on the stack; valgrind reports the read.
shoumikhin
force-pushed
the
nonzero-masked-select-device-io
branch
from
November 8, 2026 11:17
1fd85e6 to
70d8556
Compare
shoumikhin
added
release notes: backends
and removed
release notes: exir
shoumikhin
changed the title
pytorch-bot
Bot
added
the
module: vulkan
Contributor
Author
|
@claude /pr-review |
|
Claude finished @shoumikhin's task in 50m 47s —— View job PR Review: #24771Request Changes: QNN's default flow and the legacy Detailed reviewSummaryThe PR changes four things:
There is one major problem: QNN's default flow and the legacy Backward Compatibility
API Design
Testing
Performance
Code Quality
RecommendationRequest Changes.
No code was changed, and nothing was built or run. At least one reviewer read each of the 35 changed files. |
ATen returns the result of `torch.nonzero` column by column, and exported programs keep that dim order, `(1, 0)`. Most portable kernels read only the default dim order or channels last, so most ops that use the result fail at run time. For example `x[idx[:, 0], idx[:, 1]]` with `idx = torch.nonzero(x)`: ``` Check failed (all_contiguous || all_channels_last): 2 input tensors have different dim orders KernelCall failed at instruction 0:1 in operator aten::select_copy.int_out ``` A copy that keeps the layout of the result, such as `idx.float()`, `idx.clone()` or `idx / 2`, fails earlier, in `to_edge`, because the dim order ops cannot express `(1, 0)`: ``` Failed to map a given dim_order: [1, 0] to a torch.memory_format ``` `to_edge` and `to_edge_transform_and_lower` now copy the result of `nonzero` into the default dim order, before `MemoryFormatOpsPass` and before partitioning. The older `ExirExportedProgram.to_edge`, which `exir.capture` uses, does the same. So every op after it reads the default dim order, every delegate is compiled for the layout it reads at run time, and a caller that reads a returned result row by row reads it right. `as_strided` of the result still reads the result itself, because its strides refer to the layout that ATen gave the result. So a delegate that follows strides, such as OpenVINO, reads it right, and the portable `as_strided_copy`, which reads only the default dim order, still stops with an error, as before. `as_strided` of a value computed from the result, such as `idx.float()`, reads that value in the default dim order, where eager keeps the layout of `idx`. The copy moves `N * ndim` integers, which is small next to the `nonzero` scan of the whole input, but it needs its own planned memory, as large as the upper bound of the result, unless a delegate takes it into its partition. A graph without `nonzero` is left as it is, so its export does not pay for retracing the graph. Like `MemoryFormatOpsPass`, this runs only with dim order ops, which are on by default. With `_skip_dim_order=True`, and in programs exported before this change, there is no copy and the result keeps its `(1, 0)` layout. Ops that follow strides, such as `gather`, `scatter`, `unsqueeze_copy`, `narrow_copy` and `index` (as an index), read it right. Most other ops refuse it with an error, as before. The caller must follow its dim order. Exporting again without `_skip_dim_order=True` adds the copy. The new tests fail without this change, except the `as_strided` test, which fails if `as_strided` reads the copy. One runs a program that reads the result through `select_copy` and through `.float()`, and compares it with eager. It also fails without the previous fix to the `nonzero` kernel, since the copy reads the kernel output. It also checks that the copy does not bind the size that depends on the data, as `nonzero` does: a CUDA delegate that takes the copy without `nonzero` fails to lower otherwise. Another checks that a graph without `nonzero` is not retraced. Another exports through `ExirExportedProgram.to_edge` and checks that the returned result is copied and contiguous. Another checks that `as_strided` reads the result itself. A CUDA test runs a program where a CUDA `add` reads the result, another CUDA `add` reads a `mul` of it, and a portable `select` reads it, and compares it with eager. `test_optional_tensor_list` now expects the copy before `index`.
The portable `nonzero` kernel now writes its result in its declared dim order, `(1, 0)`, column by column, and export copies that result into the default dim order. The Qualcomm backend runs `RemoveRedundancy` before partitioning, and it removed every `_clone_dim_order`, so it also removed that copy. Qualcomm has no `nonzero` op, and its partitioner keeps every op that reads a size that depends on the data on the CPU, so `nonzero` and the ops that read its result run on the CPU. Without the copy, such an op, like `idx + 1` or `x[idx[:, 0], idx[:, 1]]`, fails at run time because its inputs have different dim orders, as on main. And a returned result keeps dim order `(1, 0)`, so a caller that reads it row by row, as it could before, reads it wrong. `RemoveRedundancy` now keeps a `_clone_dim_order` that reads a `nonzero` result. It still removes the others, as before. The new test fails without this change.
shoumikhin
force-pushed
the
nonzero-masked-select-device-io
branch
from
November 8, 2026 16:27
70d8556 to
cf1f16d
Compare
Contributor
Author
|
Thanks. Changes:
Not changed here:
|
Contributor
Author
|
@claude /pr-review |
|
Claude finished @shoumikhin's task in 31m 52s —— View job PR Review: #24771Approve. There are no blockers. The one minor finding is that the "Behavior change" section should warn that programs re-exported with Detailed reviewSummaryThe PR makes Backward Compatibility
RecommendationApprove Nothing blocks approval. The finding above only needs a note in the PR description.
No code was changed, and nothing was built or tested. |
Contributor
Author
|
Thanks. I reproduced it with your example: a runtime built before this change returns |
This branch has not been deployed
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.
What is wrong today
Programs declare a dim order (memory layout) per tensor, but:
nonzerokernel writes its result row by row, while the program declares it column by column (dim order(1, 0)), as PyTorch returns it. So the indices are wrong.What this change does
nonzero,gather,scatter,unsqueeze_copyandnarrow_copyfollow strides, so they read any input in its declared dim order, including the(1, 0)result of older or_skip_dim_order=Trueprograms.scatterandscatter_addwith a 0-dim index no longer read uninitialized memory.nonzeroresult into the default dim order before partitioning, also throughexir.capture, and the Qualcomm backend keeps that copy, so every op and delegate is compiled for the layout it reads.idx.float()andidx / 2now export.as_stridedof the result itself still reads the original.Behavior change
nonzeroresult row by row from older or_skip_dim_order=Trueprograms.MethodMeta, viamake_tensor_ptrwith a dim order, orTensor.fromBlobwithdimOrderon Android), or export with a contiguous example input.nonzero. Update the runtime first.What was tested
Each new test fails without its fix. On Linux x86_64, C++ kernel and util tests and exir, Vulkan, OpenVINO, WebGPU, CUDA and Python bindings tests pass, except tests that also fail on main. Apart from the callers above, programs exported on main that gave right values still do. Vulkan and WebGPU ran on a software driver, OpenVINO on its CPU plugin.
Fixes #24763