Repository navigation
Fix nonzero, masked_select, GPU constants and returned external constants, and refuse device inputs that cannot run - #24771
shoumikhin wants to merge 6 commits into
Conversation
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 new kernel test fails 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 any op that uses the result fails 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 read it as usual. 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 result. The check uses the dim order the program gets. For a 1-D input the result has a single column, and there `torch.Tensor.dim_order()` reports the default order while the program gets `(1, 0)`. The new test fails without this change, and also without the previous fix to the `nonzero` kernel, since the copy reads the kernel output.
`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. 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 host memory, which it refuses for device memory, so the exported program failed on every call with "Planned-buffer copy needs host memory on both sides". MemoryPlanningPass now checks its result: a method input on a device that got a planned buffer makes to_executorch raise a ValueError that names the input and the fix, MemoryPlanningPass(alloc_graph_input=False). The check reads the actual plan, so host inputs, planned outputs, constants kept on the device and an input that is also an unplanned output are accepted. It only runs with device-aware memory planning, which the skip flags require. Tests that exported this layout only to inspect the graph now set alloc_graph_input=False and no longer expect a planned input buffer. The config docstring and the CUDA and TensorRT pages now say the export raises. Test plan: test_skipping_copies_requires_unallocated_graph_io expects the ValueError and fails without the change. Ran the propagate device pass, memory planning, program, emitter and pass tests.
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: there is no operator for the CUDA partitioner to take. to_executorch() then raised an InternalError (and crashed before pytorch#24202). A program stores constants as host bytes whatever device they came from. So the emitter now reads a constant whose storage is not on the CPU 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. Storage with no data, such as "meta", still raises an InternalError. This also lets a GPU module lowered with no backend at all export, which pytorch#24202 made raise: it now gives the same program as exporting the module from the CPU. Test plan: a CUDA test exports a GPU module that returns (x * 2, buffer) with the CudaPartitioner, runs it, and compares with eager. It fails without the change. The meta test is renamed and matches the new error.
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 header documentation now says that constants count as memory planned. This changes 2 things that callers can see, 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 from a program exported with external constants, and a C++ test that checks the external constant attributes of an existing test program report is_memory_planned().
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. ❌ 3 Pending, 2 Unrelated Failures, 1 Unclassified FailureAs of commit abe9f62 with merge base 75af78c ( 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:
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. |
What is wrong today
torch.nonzeroreturns wrong indices with no error. PyTorch lays its result out column by column, and the portable kernel writes row by row. Ops that read the result, likex[idx[:, 0], idx[:, 1]], also fail at run time.to_edgerejects a directtorch.masked_selectcall as "not in Core ATen opset", although a portable kernel exists.skip_h2d_for_method_inputs=True, memory planning still reserves a GPU buffer by default for each input left on the GPU. The runtime cannot copy the caller's GPU tensor into it, so every call fails.to_executorch()raise "Constant tensor storage is on device 'cuda:0'".external_constants=Truereturns zeros, and later reads of that constant see zeros too.What each commit does
nonzerokernel writes each index through the strides of its output, so it is right for both layouts.LegalizePortableDimOrderPasscopies thenonzeroresult into the default layout, so later ops and callers read it as usual. The copy needs memory as large as the result.aten.masked_select. Inputs with 2 elements or fewer still fail into_executorch(), and an input that broadcasts to a bigger shape can fail at run time, both with an error.MemoryPlanningPassrefuses a planned buffer for a method input left on a device. Such a program used to export and fail on every call; nowto_executorch()raises aValueErrorthat names the input and the fix,MemoryPlanningPass(alloc_graph_input=False).MethodMetacounts a constant stored in the.ptdfile as memory planned, like an inline constant. So the bindings no longer give it an empty output buffer,is_memory_planned()is true for it, andset_output_data_ptron such an output returnsInvalidStateinstead of overwriting the constant.What was tested
Each commit adds a test that fails without it. The kernel,
exir, emitter, program,MethodMeta, pybindings and CUDA backend tests pass on Linux x86_64 with an H100, except failures that main also has.Fixes #24763
Fixes #24764
Related: #21580, #21545