Repository navigation
Fix export of zero-stride constants and inputs, and of constants that delegates share - #24404
shoumikhin wants to merge 12 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 f7e48fa 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
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.
There was a problem hiding this comment.
Agreed on both. The warning should name the reason that applied, and the ops or tags outside the delegate (or the returned output). The docstring should say why each condition exists. Both change only text and logging, so I will do them in a follow-up PR rather than restart review here.
| # 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.
There was a problem hiding this comment.
I ran this at the PR's head with an input of strides (0, 1) and y = x.clone(); return y * 2, x.view(1, 4) + 1:
- Default
to_edge: the clone gets(4, 1). _skip_dim_order=True: the clone keeps(0, 1)in its metadata, as you said. Still,to_executorchworks and both outputs match eager.
So this case works today. I will add the clone and _skip_dim_order=True cases to test_emit_user_input_with_zero_stride in the follow-up PR.
| 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.
There was a problem hiding this comment.
Right. This is already the case on main, and not only for zero strides. On main, convert_constant_dim_order_pass keeps only contiguous and 4-D channels-last constants, so a 5-D channels_last_3d constant becomes contiguous. I checked a channels_last_3d constant with a zero stride on a size-1 dim, added to a channels_last_3d input: the add fails its dim order check at run time on main and with this PR. The 4-D version fails to export on main and works with this PR. I will handle channels_last_3d in both places in a follow-up.
| 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.
There was a problem hiding this comment.
Agreed. A change that makes one copy per tag, with all untagged users sharing one, is already written as its own commit. It goes into #24317, after this PR merges.
| 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?
There was a problem hiding this comment.
Done. The PR description now has a "Behavior change" paragraph that names this case. With the reader delegated, that model failed to lower on main (Buffer output x does not point to a buffer that exists). Now it lowers and runs, so no working program changes here.
|
Thanks for the careful review.
The Arm TOSA copy goes to an Arm PR, as you suggested. |
PyTorch ignores the stride of a size-1 dim when it checks whether a tensor is contiguous. ExecuTorch derives a tensor's dim order (the order of its dims in memory) from all of its strides. So PyTorch can call a constant contiguous while its strides give a non-default dim order. For example, a buffer set to `torch.randn(1, 8, 4).transpose(0, 1)` has sizes [8, 1, 4] and strides (4, 32, 1). PyTorch calls that contiguous, so convert_constant_dim_order_pass left it alone, and the program stored it with dim order (1, 0, 2). The same happened when constant_prop_pass folded `self.queries.transpose(0, 1) * 0.5` into a constant. Ops that read the constant then failed at run time: - portable bmm: "3 input tensors have different dim orders" - XNNPACK: the delegate reorders the input sizes by the dim order, gets [1, 8, 4] instead of [8, 1, 4], and fails with "Propagating input shapes failed" Now convert_constant_dim_order_pass also checks the dim order that a constant's strides give, after it removes zero strides. If that order is neither the default nor channels-last, it copies the constant, buffer or parameter into the layout PyTorch reports, contiguous or channels-last. The data keeps its order in memory, only the strides change. A constant whose strides already give the default or channels-last order is left as it is, even when PyTorch would call it both. For example, the weight of a channels-last Conv2d with one input channel keeps its channels-last dim order. The new check does not apply to user inputs. Test: test_emit_constants_with_size_one_dims exports a model with these buffers, runs it with the portable kernels, and compares with eager: - a transposed [1, 8, 4] buffer read by bmm, which failed before; - a [1, 3, 4, 4] buffer with strides (5, 1, 12, 3), channels-last except for its size-1 batch dim, added to a channels-last input, which also failed before; - a channels-last [2, 1, 4, 4] buffer added to a channels-last input, which must keep its channels-last dim order.
|
@claude /pr-review |
|
Claude finished @shoumikhin's task in 37m 24s —— View job Approve: four minor findings, none a regression. The size-1 dim-order fix doesn't reach method inputs or constants created after PR Review: #24404SummaryThe PR copies constants that have zero strides, or size-1-dim strides that give a non-default dim order, and removes zero strides from method inputs. It also keeps mutated buffers and shared constants correct when a delegate covers only some of their users. There are four minor findings and none is a regression. Python could not run in this job, so all findings come from reading the code. Code Quality
Testing
RecommendationApprove Nothing blocks merging, and nothing is a regression. The two tests are small to add. The input and emitter gaps can go on the follow-up list.
|
|
Thanks for the review.
|
What is wrong
Zero strides. ExecuTorch cannot store a zero stride, and
contiguous()keeps one on a size-1 dimension:So export fails with
0 in strides is not supported for ExecuTorch.whenconstant_prop_passfoldsself.scale.expand(1), a parameter, buffer or constant method holds such a tensor, or an example input is expanded.Strides of size-1 dimensions. PyTorch also ignores the stride of a size-1 dimension when it checks contiguity, but ExecuTorch derives the dim order (the order of dimensions in memory) from every stride:
An op that reads such a constant fails at run time: portable
bmmwith3 input tensors have different dim orders, XNNPACK withPropagating input shapes failed. This happens for a parameter or buffer likeq, and whenconstant_prop_passfoldsself.queries.transpose(0, 1) * 0.5.Constants that ops in different delegates read.
duplicate_constant_nodegives each such op its own copy.cat, kept the original, and lowering failed.tag_mutated_buffergave a mutated buffer to a delegate even when some readers or writers stayed outside. A reader outside got a stale copy, or lowering failed (new value computed outside, read outside, or also returned).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 one from method inputs, so an input with a zero stride on a dimension larger than 1 must be passed contiguous.convert_constant_dim_order_passalso checks the dim order that a constant's strides give. If it is neither the default nor channels-last, the pass copies the constant into the layout PyTorch reports. Constants whose strides already give one of those orders do not change, so a channels-last weight with a size-1 channel stays channels-last.duplicate_constant_nodecopies the constant for every op that reads it, including from a list, and 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 it stays in the owning program and loses any tag, with a warning.
Behavior change. A mutated buffer used by ops in different partitions, which used to give stale reads, now stays in the owning program and is passed to each delegate: correct, but copied on every call. A buffer whose new value is a graph input, as in
self.state.copy_(x), is no longer tagged, so an existing test intest_partitioner.pynow expects no tag. With the reader delegated, that model failed to lower on main, and now it lowers and runs.Testing
exir/emit,exir/testsandexir/backend/testeach fail without their fix.qabove give identical.ptefiles.Tested on Linux with torch 2.14.0 (CPU).