Repository navigation
Fix nonzero, masked_select, GPU constants and returned external constants, and leave device inputs unplanned - #24771
shoumikhin wants to merge 6 commits into
Conversation
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 Pending, 2 Unrelated Failures, 2 Unclassified FailuresAs of commit ee9b608 with merge base a84676e ( UNCLASSIFIED FAILURES - DrCI could not classify the following jobs because the workflow did not run on the merge base. The failures may be pre-existing on trunk or introduced by this PR:
FLAKY - The following job failed but was likely due to flakiness present on trunk:
BROKEN TRUNK - The following job failed but was 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
|
@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. |
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 portable `gather` and `scatter` kernels read their index tensor as if it were stored row by row. With a `nonzero` result as the index, they were right only because of the old row-by-row writes. They now read each index value through the strides of the index tensor, as `index` already does, so a program where `nonzero` feeds `gather` or `scatter` keeps working, including programs exported before this change. A `nonzero` result now has the `(1, 0)` layout that the program declares. 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 backend does. The new kernel tests fail without the fix.
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 ``` `LegalizePortableDimOrderPass` already inserts dim order copies for portable kernels during `to_executorch`. It now also copies the result of `nonzero` into the default dim order, so the ops after it, and the caller, read it as usual. A delegate that reads the result still gets the result itself, because the delegate was compiled for its dim order. 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. Like the other copies of this pass, this needs 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: `gather`, `scatter` and `index` read it right, the other ops refuse it with an error as before, and the caller must follow its dim order. Exporting again without `_skip_dim_order=True` adds the copy. The new tests fail without this change. One runs a program and compares it with eager; it also fails without the previous fix to the `nonzero` kernel, since the copy reads the kernel output. The other checks that a delegate reads the result itself while a portable op reads the copy.
`to_edge` rejects a direct `torch.masked_select` call: ``` SpecViolationError: Operator torch._ops.aten.masked_select.default is not in Core ATen opset ``` `masked_select` is not a core ATen op, and PyTorch has no decomposition for it, so it stays in the graph. ExecuTorch already has a portable kernel for `masked_select.out`, so the op can run as it is. Plain `x[x > 0]` is not affected, because export turns it into `aten.index.Tensor`, but tools that emit `masked_select` directly hit this error. This lists the op as an exception in the edge verifier, as the antialiased resize ops are. Limits remain that come from the output size bound that PyTorch's fake `masked_select` sets during `torch.export`. A decomposition in ExecuTorch runs after that, so it does not remove them (it fixes only inputs of exactly 2 elements): - The bound is the size of the input, not of the input broadcast to the mask. If the input broadcasts to a bigger mask, a call that selects more elements than the input has fails at run time. - An input with 2 elements or fewer gets no bound, so `to_executorch` fails with an error about an unbounded size. These limits stop with an error, never with wrong values. Boolean indexing such as `x[mask]` (after `torch.broadcast_tensors` if the shapes differ) avoids them, because export turns it into `aten.index.Tensor`. The new `to_edge` test fails without the change.
With skip_h2d_for_method_inputs, a method input that a device delegate reads stays on the device, and the caller passes their own device tensor. Memory planning still reserved a buffer for that input by default. The runtime fills a planned input by copying into it, which it refuses to do for device memory, so the exported program failed on every call with "Planned-buffer copy needs host memory on both sides". MemoryPlanningPass(alloc_graph_input=False) avoided that, but it also unplans every host input of the method. PropagateDevicePass now marks such an input, as it marks constants that it keeps on the device, and memory planning never gives a marked input a buffer, whatever alloc_graph_input says. The method then reads the caller's tensor in place. Host inputs still follow alloc_graph_input, and the memory planning verifier leaves marked inputs out of its check. The PropagateDeviceConfig docstring and the CUDA and TensorRT pages no longer ask for alloc_graph_input=False. Tests: test_inputs_left_on_device_are_not_planned replaces a test that pinned the old default, and a CUDA test now runs such a program with the default planning. Both fail without the change. Assertions that expected a planned buffer for such an input are gone, since that plan fails on every call.
abe9f62 to
ad64232
Compare
shoumikhin
changed the title
|
@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. |
When a module is still on the GPU at export, as in flows that compile it there first such as Torch-TensorRT, a constant that no delegate takes stays in the program with its storage on the GPU. That happens when forward() returns a buffer, alone or next to a computed result, or when a delegate reads a constant from outside it. to_executorch() then raised an InternalError (and crashed before pytorch#24202). A program stores constants as host bytes whatever device they came from. So when only delegates, copies between host and device memory, or the method output read such a constant, the emitter now reads it from a host copy instead of raising. The constant keeps its place and device information in the program, so the runtime treats it as it would the same constant exported from the CPU. The InternalError stays when an operator outside any delegate reads the constant. That operator was traced on the GPU but runs on a CPU kernel, which may not match (pytorch#21545): a GPU batch norm lowered with no backend, for example, gets statistics of a shape the CPU kernel rejects, and fails on every call. Storage with no data, such as "meta", still raises too. Test plan: CUDA tests export a GPU module that returns (x * 2, buffer) with the CudaPartitioner, and a GPU module whose constant a delegate reads from outside it. Both run, match eager, and fail without the change. Another test checks that a GPU Linear lowered with no backend still raises, and the "meta" test now also covers a buffer that only the output reads.
A model that returns one of its constants (a buffer or a parameter) gives zeros at run time when it is exported with external_constants=True. There is no error. The same model with its constants kept inside the .pte file returns the right values. MethodMeta decides whether an output is memory planned by looking for an allocation or an inline data buffer. A constant stored in a .ptd file has neither, so it looks like an output with no storage. The Python bindings then give that output a new zero-filled buffer, and set_output_data_ptr replaces the constant's data pointer with it. From then on every read of the constant sees zeros, including operators that use it to compute other outputs. Treat a tensor whose data lives in an external file as memory planned, the same way inline constants already are. One helper now computes this for inputs, outputs and attributes, and the C++ and Objective-C header documentation now says that constants count as memory planned. Callers can see this change for external constants only. Inline constants already behaved this way: - TensorInfo::is_memory_planned() now returns true for an output or an attribute that is an external constant. The Python, Swift and JavaScript bindings show this value. - Method::set_output_data_ptr and Module::set_output on such an output now return InvalidState, instead of silently replacing the constant's data. Add a Python test that returns a buffer, and a sum that uses it, from a program exported with external constants. It checks that the buffer output reports is_memory_planned() and that both outputs stay right over 2 runs. A C++ test checks that the external constant attributes of an existing test program report is_memory_planned().
|
@claude /pr-review |
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. |
What is wrong today
torch.nonzerosilently returns wrong indices: PyTorch lays the result out column by column, but the portable kernel writes row by row. Ops likex[idx[:, 0], idx[:, 1]]then fail.to_edgerejectstorch.masked_selectas "not in Core ATen opset", although a portable kernel exists.skip_h2d_for_method_inputs=Trueand default memory planning, inputs left on the GPU get buffers the runtime cannot fill, so every call fails.to_executorch()raise "Constant tensor storage is on device 'cuda:0'".external_constants=True, Python gets zeros for a returned constant, and later reads of it see zeros too.What each commit does
nonzerokernel writes through its output strides.gatherandscatterread their index through its strides too, so old programs wherenonzerofeeds them still work.to_executorchcopies thenonzeroresult into the default layout for later portable ops and callers. Delegates get the result column-major, as declared, and must follow its dim order.aten.masked_select. PyTorch's export bounds its output size wrongly: an input that broadcasts to a bigger mask can fail at run time, and inputs of 2 elements or fewer fail into_executorch(), both with an error. Same-shapex[mask]avoids this.alloc_graph_input.MethodMetacounts a.ptdconstant as memory planned, like an inline one. The bindings stop giving it a zero-filled output buffer, andset_output_data_ptron it returnsInvalidState.For C++ callers: programs exported before this change or with
_skip_dim_order=Truereturnnonzeroresults column-major, as declared. Read them through their dim order, or re-export without_skip_dim_order.What was tested
Each commit adds tests that fail without it. Tests for the changed code (kernels, method,
exir, pybindings, CUDA) pass on Linux x86_64 with an H100, except failures also seen without this change. Old programs feedingnonzerotogatherorscattermatch eager.Fixes #24763
Fixes #24764
Fixes #21580
Related: #21545