Repository navigation
Fix export of zero-stride constants and inputs, and of constants that delegates share - #24404
shoumikhin wants to merge 11 commits into
Conversation
constant_prop_pass folds ops whose inputs are all constant into new
constant tensors. ExecuTorch cannot store a tensor with a zero stride, so
the pass copies a folded tensor that has one. It used contiguous() for
that copy, but PyTorch ignores the stride of a size-1 dim (and every
stride of an empty tensor) when it decides whether a tensor is
contiguous, so contiguous() can return the tensor unchanged. For example,
expand(1) of a 0-d tensor:
t = torch.tensor(2.0).expand(1)
t.stride() # (0,)
t.is_contiguous() # True
t.contiguous().stride() # still (0,)
The zero stride then reached to_executorch, which failed with
"0 in strides is not supported for ExecuTorch."
Copy with clone(memory_format=torch.contiguous_format) instead, which
always gives standard strides. Folded constants without a zero stride are
not touched, and the ones contiguous() already copied come out the same.
shoumikhin
added
the
release notes: exir
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/24404
Note: Links to docs will display an error until the docs builds have been completed. ❌ 2 New Failures, 1 Unclassified FailureAs of commit 0b16218 with merge base 05092c1 ( NEW FAILURES - The following jobs have failed:
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:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
A parameter, buffer, or constant method value made with expand() can
have a zero stride on a dim of size 1, for example:
t = torch.tensor(2.0).expand(1)
t.stride() # (0,)
t.is_contiguous() # True
t.contiguous().stride() # still (0,)
PyTorch ignores the stride of a size-1 dim, and every stride of an empty
tensor, when it decides whether a tensor is contiguous. So the code that
gives constants a layout ExecuTorch can store left these tensors as they
were, and export failed with "0 in strides is not supported for
ExecuTorch." A parameter or buffer failed in to_edge, and a constant
method value failed in to_executorch.
Add without_zero_strides() to exir/tensor.py. It returns a contiguous
copy of a tensor that has a zero stride, and the tensor itself
otherwise. constant_prop_pass uses it in place of its own copy of the
same check, and TensorSpec.from_tensor(const=True), which constant
methods use, uses it after contiguous().
convert_constant_dim_order_pass now also copies parameters, buffers and
lifted constants that have a zero stride. It keeps a channels-last
tensor channels-last, as it does for every other constant, and does not
log its "neither contiguous nor channels-last" warning for these
tensors, since their layout does not change.
Constants without a zero stride come out exactly as before.
shoumikhin
changed the title
test_emit.py now imports convert_constant_dim_order_pass directly. The emit target skips autodeps, so list the dep by hand.
|
@claude /pr-review |
|
Claude finished @shoumikhin's task in 24m 48s —— View job No blockers. There are three minor findings: Detailed reviewPR Review: #24404SummaryThe PR adds Code Quality
Testing
RecommendationApprove Nothing blocks this PR. All three findings are minor and none is a regression. The channels-last fix and the |
When a constant has users in different partitions, or in a partition and on the CPU, to_backend gives each user its own copy of the placeholder. duplicate_constant_node only swapped the copy into the top-level args and kwargs of a user, so an op that reads the constant from a list, such as cat or the indices of index_put, kept the original. Lowering then failed with "constant data node ... is tagged with ... but has user ... which has tag (None)". For example, a Core ML model whose index buffer is read by an index_put in Core ML and by another index_put on the CPU failed this way. The pass now swaps the copy in wherever the user reads the constant. A mutated buffer keeps its list readers on the original, since a copy would not see the buffer's updates, so such a model still fails to lower instead of reading a stale value.
shoumikhin
changed the title
|
@claude /pr-review |
|
Claude finished @shoumikhin's task in 24m 15s —— View job One major finding: Detailed reviewPR Review: #24404SummaryThis PR does two things:
There is one major finding, in how Code Quality
Testing
RecommendationRequest Changes
|
A mutable buffer, such as a KV cache, changes on every call. When some of the ops that read or write it go to a delegate and others stay in the owning program, tag_mutated_buffer used to tag the buffer for the delegate anyway, because it ignored untagged readers and never looked at the op that computes the new value. - If another op read the buffer outside the delegate, to_backend gave that op its own copy. A copy keeps the value from export time, so the program lowered and gave wrong results from the second call on. For example, with a buffer b = zeros(2), forward doing y = x * b, z = x + b, b.add_(1), and only the multiply delegated, the second call returned [1, 1] instead of [3, 3]. - If the new value was computed outside the delegate, lowering failed with "Buffer output ... does not point to a buffer that exists". tag_mutated_buffer now tags a mutated buffer only when all of its readers and writers are in the same partition, as its docstring already said. Otherwise the buffer stays in the owning program and the delegates read it as an input. duplicate_constant_node now refuses to copy a mutated buffer, for direct and list readers alike, with an error that says how to tag it, instead of making a copy that never sees the updates. The existing test_tag_mutated_buffer_detects_indirect_buffer_mutation computed the new value from a graph input outside the partition, where the tag it expected made to_backend fail with "Buffer output x does not point to a buffer that exists". It now computes the value inside the partition.
…ode copies When the graph output also returned a constant that ops in two partitions read, duplicate_constant_node made a copy for the output, which nothing then read, and left the original with both the first op and the output. Lowering then failed with "has user (output) which has tag (None)", for example for a model that returns torch.cat([x, const]), const with only the cat delegated. Now the output keeps the original, as the output specs name it, and every other user reads its own copy. Swapping the constant in for each user takes one replace_input_with call, which also reaches a constant inside a list, so the separate argument and keyword rewrites are gone. A to_backend test now lowers a constant that cat ops in two partitions read, with and without the output also returning it, and checks which placeholder each delegate and the output read. It fails without either the list fix or this one, so it replaces the pass-level list test, which checked the same cases.
without_zero_strides always copied to the default layout, so convert_constant_dim_order_pass had its own branch to keep a channels-last constant channels-last, and constant_prop_pass, which calls the helper directly, did not. A channels-last tensor with a zero stride, such as w.expand(1, 3, 4, 4) for a w with strides (1, 12, 3) and size (3, 4, 4), was then stored channels-last as a buffer but contiguous when folded. The helper now keeps a channels-last tensor channels-last, so both passes follow one rule and the pass loses its extra branch. TensorSpec.from_tensor is unchanged, since it makes the tensor contiguous first.
test_emit_folded_constant_with_zero_stride also ran the folded program on the portable runtime. That part passes even without the constant_prop_pass fix, because to_edge now removes the zero stride itself, and test_emit_user_constants_with_zero_stride already runs zero-stride constants end to end. The test now only checks the stride of the folded constant, which is what fails without the fix. It is renamed to test_constant_prop_pass_folds_without_a_zero_stride, since it no longer emits anything. It stays in test_emit.py because exir/tests/test_passes.py is not collected by the open source test run.
shoumikhin
changed the title
|
@claude /pr-review |
|
Claude finished @shoumikhin's task in 27m 35s —— View job One major finding: Detailed reviewPR Review: #24404SummaryThe PR does three things:
There is one major finding, in Backward Compatibility
API Design
Testing
Code Quality
RecommendationRequest Changes The write-only buffer change alters main's behavior in a util that eight in-tree partitioners call. By reading, it can break lowering. The fix is one condition plus a test. The other three findings are minor, and the Arm cleanup can be a follow-up. |
…wn it tag_mutated_buffer counts the op that computes a buffer's new value as a user. So a buffer that no node reads got that op's tag, and the delegate owned it. That fails in lowering when the new value is also a user output, for example `self.state.copy_(x + x); return self.state`: "Constant state was tagged to be consumed by the buffer ...". The same failure happens when the delegate also reads the buffer, as in `self.state.add_(x); return self.state`, and that one fails on main too. Now a buffer is tagged only when every op that reads it (its old value or its new value) or writes it has the same tag, and no writer is a user output. Before, an op that read the new value was not counted, so `self.state.add_(x); return torch.sin(self.state)` with only the add delegated failed in to_backend, on main too. A buffer that no node reads stays in the owning program, as on main, unless the partitioner already gave it to the delegate that writes it (QNN tags such buffers itself). A buffer that stays also loses any tag the partitioner put on it before, and a warning says so. Tests: - test_to_backend_keeps_a_mutated_buffer_out_of_a_partial_delegate adds a case with the new value read outside the delegate, a case with the new value returned, write-only cases with and without a tag from the partitioner and with the new value returned, and a case with the readers in two delegates. It also runs the case with a delegate that reads the buffer twice, and compares with eager. - test_tag_mutated_buffer_detects_indirect_buffer_mutation keeps its copy_(x) model, where the new value is a graph input, and checks the buffer stays untagged.
An example input can have a zero stride on a size-1 dim, for example after an expand, and PyTorch still calls it contiguous. to_executorch then failed with "0 in strides is not supported for ExecuTorch." when it built the spec of that input. convert_constant_dim_order_pass, which to_edge runs, already removes zero strides from constants. It now does the same for every other placeholder. An input with a zero stride becomes contiguous, or stays channels-last, as constants do. Other inputs do not change. Test: test_emit_user_input_with_zero_stride exports with an input of stride (0, 1), runs the program, and compares with eager.
shoumikhin
changed the title
|
@claude /pr-review |
|
Claude finished @shoumikhin's task in 33m 7s —— View job Approve with five minor findings. The most notable: AOTI's fallback re-tags the mutated buffers that Detailed reviewPR Review: #24404SummaryThis PR lets ExecuTorch export constants and user inputs that have a zero stride. It also makes Code Quality
Testing
API Design
RecommendationApprove Nothing blocks merging. The gaps in the backends and in
|
The Windows unit test job runs pytest against a build that does not register the demo backend, so the runtime part of the new mutated buffer test fails there with "Backend BackendWithCompilerDemo is not registered". Deselect it in pytest-windows.ini, next to the other tests that run the demo backend.
|
@claude /pr-review |
|
Reviewing PR #24404
|
|
Thanks for the review. Follow-ups. These are fair. I will do them in a separate PR:
VGF already calls Two buffers with one writer, and a buffer placeholder as the writer. Guards from the zero-stride check. With an unbacked size |
|
@claude /pr-review |
|
Claude finished @shoumikhin's task in 24m 33s —— View job Approve with two minor findings: Detailed reviewPR Review: #24404SummaryThe PR makes constants and inputs with a zero stride exportable. It also has Code Quality
Testing
RecommendationApprove Neither finding blocks merging. The
|
|
Thanks. Both points are fair. Neither one breaks anything that works on main today, so I will do them in the same follow-up PR as the items in my earlier reply:
|
|
Why are these 2 conditions? its new value is not returned; neither obviously have correctness issues to me |
|
Good question. Short answer: the first condition is needed for lowering to work. The second is not a correctness fix. It keeps main's behavior. "Its new value is not returned." If a delegate owns a buffer and the op that writes it is also a user output, "Some op reads it, unless the partitioner already tagged it for its writer." You are right, this one is not about correctness. On main, Dropping it would move write-only buffers into the delegate that writes them, whenever their writer is tagged. In the test, that still lowers. But it changes behavior for those backends, so I would keep it out of this fix. |
metascroy
left a comment
metascroy
left a comment
There was a problem hiding this comment.
Thanks, the fixes look correct and the tests are thorough. Most of my comments are inline. Two general ones:
1. The tag_mutated_buffer change reaches every backend that uses it. On main it's called by Core ML, Core AI, QNN, Vulkan, MLX, Arm VGF, AOTI and native. Before this change, a buffer used by more than one partition went to min(user_tags), and readers outside that delegate got a stale copy. Now it stays in the owning program and is passed to the delegates as an input and output. That's correct, but it can be much slower: a KV cache that used to be delegate state (for example Core ML MLState) is now copied in and out on every call. Since this needs sign-off from those backends' owners and the zero-stride fix is low risk, I'd consider splitting the two.
2. This affects the Core ML change in #24317. That PR adds _check_state_is_used_in_one_place, which raises when Core ML state is used in more than one place, because "each place would use its own copy." With this PR the buffer stays in the program instead, so the stated reason no longer holds. #24317 should decide whether to keep raising (the user did ask Core ML to own the state), warn instead, or rely on this behavior.
Minor: the Arm TOSA partitioner keeps its own copy of the old rule, so for now Arm behaves differently from the shared helper. Deferring that to an Arm PR sounds fine.
The three failing checks look unrelated: the NXP job couldn't fetch a git ref, the Voxtral job got a 429 from Hugging Face, and unittest-cuda failed only test_antialiased_resize, on an Inductor make_fallback(aten.arange) assert.
| continue | ||
| had_tag = node.meta.pop("delegation_tag", None) is not None | ||
| if had_tag or (node.users and user_tags != {None}): | ||
| logging.warning( |
There was a problem hiding this comment.
This warning is the only sign that a backend has lost ownership of its state, which can be a big slowdown. It lists three possible reasons without saying which one applied. Could it name the actual reason and the ops/tags outside the delegate (or the returned output)? That would make the fallback much easier to diagnose.
Also, the answer you gave Jacob on the two extra conditions (one keeps lowering from hitting an assert, the other keeps main's behavior) would be worth putting in the docstring.
| # store either. Constants are already normalized, so this only changes inputs. | ||
| for node in exported_program.graph.nodes: | ||
| if node.op == "placeholder" and isinstance(node.meta.get("val"), torch.Tensor): | ||
| node.meta["val"] = without_zero_strides(node.meta["val"]) |
There was a problem hiding this comment.
This rewrites only the placeholder's meta["val"]. With plain torch, exporting a model whose input has strides (0, 1) leaves clone and alias of that input with stride (0, 1) in their metadata (mul gets (4, 1)).
On the default path the later passes in to_edge probably recompute that metadata from this placeholder. With _skip_dim_order=True, though, this pass runs last, and the comment in _program.py says ops computed from the normalized tensors keep the old strides.
Could you add x.clone() and a _skip_dim_order=True case to test_emit_user_input_with_zero_stride? I haven't run it with this PR applied, so this may be fine.
| if not tensor.is_contiguous() and tensor.is_contiguous( | ||
| memory_format=torch.channels_last | ||
| ): | ||
| return tensor.clone(memory_format=torch.channels_last) |
There was a problem hiding this comment.
Nit: this keeps channels-last only for 4-D tensors. A 5-D channels_last_3d tensor with a zero stride is copied as contiguous, which changes its dim order. That's rare, but it may be worth handling or noting in the docstring.
| users = [user for user in node.users if user.op != "output"] | ||
| if len(users) == len(node.users): | ||
| users = users[1:] | ||
| for ith, user in enumerate(users): |
There was a problem hiding this comment.
Nit, and already true on main: this still makes one copy per user, even for users that share a tag. Copying once per tag would mean fewer placeholders. Fine to leave for later.
| delegate_tag = "test_partition" | ||
| for user in state_node.users: | ||
| user.meta["delegation_tag"] = delegate_tag | ||
| for compute_new_value, expected_tag in ((True, delegate_tag), (False, None)): |
There was a problem hiding this comment.
The original case here (copy_(x), where the new value is a graph input) used to expect the buffer to be tagged and now expects no tag. The reasoning makes sense, but since this reverses an existing assertion, could you mention it in the PR summary as a behavior change?
What is wrong
Zero strides. ExecuTorch cannot store a zero stride, and PyTorch's contiguity check ignores the stride of a size-1 dimension:
So export fails with
0 in strides is not supported for ExecuTorch.in these cases:constant_prop_passfoldsself.scale.expand(1).Constants that ops in different delegates read.
to_backendgives each such op its own copy (duplicate_constant_node).cat, kept the original. Lowering then failed withconstant data node ... has user ... which has tag ....tag_mutated_buffergave a mutated buffer to a delegate even when some readers or writers stayed outside. A reader outside got a stale copy.self.state.copy_(x + x); return self.state, went to the delegate that writes it, and lowering failed.What this change does
without_zero_strides()inexir/tensor.pycopies a tensor with a zero stride, keeping channels-last.constant_prop_pass,convert_constant_dim_order_pass(run byto_edge) andTensorSpec.from_tensor(const=True)use it. That pass also removes a zero stride from method inputs. An input whose zero stride is on a dimension larger than 1 must then be passed contiguous.duplicate_constant_nodecopies the constant for every op that reads it, including from a list. The graph output keeps the original. It raises instead of copying a mutated buffer.tag_mutated_buffergives a mutated buffer to a delegate only when these hold:Otherwise the buffer stays in the owning program and loses any tag, with a warning.
Testing
test_emit.pyandtest_tensor.py, and intest_passes.pyandtest_utils.pyunderexir/backend/test. Each fails without its fix. The mutated buffer test runs one program twice against eager.exir/emit,exir/testsandexir/backend/testpass as before..ptefiles.Tested on Linux with torch 2.14.0 (CPU).