Repository navigation
exir and MLX: fix folded-output clones, shared constants, and constant boolean masks - #24481
shoumikhin wants to merge 4 commits into
Conversation
shoumikhin
added
the
release notes: exir
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/24481
Note: Links to docs will display an error until the docs builds have been completed. ⏳ No Failures, 162 PendingAs of commit 5e7f110 with merge base a489dca ( 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 19m 19s —— View job Summary: Nothing blocks merging. The findings are three test gaps (the new per-tag grouping in Detailed reviewPR Review: #24481SummaryThe PR does two things:
Nothing blocks merging. The issues below are test gaps, one note for direct callers, and one optional simplification. Code Quality
Testing
Backward Compatibility
RecommendationApprove. Nothing blocks. The tests above are worth adding in this PR, most of all the case where a CPU op is the first user, because that is the path that re-tags the original placeholder. I did not run any tests: this job can't run Python. Everything above comes from reading the code. |
ef572cf to
f4ee3c0
Compare
shoumikhin
changed the title
f4ee3c0 to
460b697
Compare
|
Thanks for the review. The latest push changes 2 tests:
What I left as is:
About direct callers: a graph with no delegation tags now gets no copies. The docstring, the commit message and the PR description say this, so it will show up in the release notes. The only non-test caller in the tree runs after tagging. |
|
@claude /pr-review |
|
Claude finished @shoumikhin's task in 17m 26s —— View job Summary: Nothing blocks merging. There are three minor findings: one subtest of the new Detailed reviewPR Review: #24481SummaryThe PR makes two changes:
Nothing blocks merging. The findings are two test gaps and an Edge-only default in the same pass. That default predates the PR. Code Quality
Testing
RecommendationApprove. Nothing blocks. The I did not run any tests, because this job can't run Python. Everything above comes from reading the code. |
460b697 to
82171cc
Compare
|
Thanks, good catch on the test. The latest push fixes it:
What I left as is:
|
|
@claude /pr-review |
|
Claude finished @shoumikhin's task in 16m 19s —— View job Summary: Nothing blocks merging. There are two minor findings: no test has one Detailed reviewPR Review: #24481SummaryThe PR makes Testing
Backward Compatibility
RecommendationApprove. Nothing blocks. The test case is a few lines in the existing test. The dialect risk is latent, so a docstring line is enough. I did not run any tests, because this job can't run Python. Everything above comes from reading the code. |
82171cc to
1f5b7dd
Compare
|
Thanks. Both points are fixed in the latest push:
|
|
@claude /pr-review |
|
Claude finished @shoumikhin's task in 14m 9s —— View job Summary: This round found no new issues, and nothing blocks merging. Detailed reviewPR Review: #24481SummaryThe PR does two things:
This round found no new issues. Two points from earlier rounds came up again and are left out: choosing the clone by RecommendationApprove. Nothing blocks. I did not run any tests, because this job can't run Python. Everything above comes from reading the code. |
|
Claude finished @shoumikhin's task in 24m 41s —— View job Summary: Nothing blocks merging. There are two minor findings: an MLX builder comment still says shared constants are copied once per user, and the new constant-mask lowering makes one int64 index per mask dimension, so a multi-dimensional mask goes through a gather path that refuses dynamic dimensions. Detailed reviewPR Review: #24481SummaryThe PR makes three changes:
Nothing blocks merging. Both findings are minor, and both involve the MLX side. The mask handler ( Code Quality
Performance
RecommendationApprove. Nothing blocks. The comment is a quick fix. The mask layout can wait unless the target models have large multi-dimensional masks or dynamic shapes. I did not run any tests, because this job can't run Python. Everything above comes from reading the code. |
052ecdd to
e6ca3a1
Compare
|
Thanks for the review. Both points: Builder comment. You are right that the grouping made it stale. I checked whether anything still gives one delegate two placeholders of the same constant. With the grouping, the shared-slot branch in the MLX builder never runs: not on the MLX partitioner tests, and not on graphs that read a constant several times inside and across delegates. Without the grouping it does run. So the branch and its comment are removed in the grouping commit. Mask layout. I kept one index per mask dimension here. A multi-dimensional mask with a dynamic dimension outside the mask is refused cleanly, so the index runs on the CPU. Static multi-dimensional masks lower through the gather path, and the tests run them. Merging the masked dimensions first would make the index handler emit its own reshape and gather, which is a bigger change than this one. The masks that motivated this are small and static, so the extra position memory does not matter for them. The new head also adds a commit that copies the shared constants of all delegates in one pass. It recompiles and validates once, not once per constant. |
|
@claude /pr-review |
|
Claude finished @shoumikhin's task in 23m 34s —— View job Summary: Nothing blocks merging, with two minor findings: the original placeholder does not go to the first reader in graph order as the e6ca3a1 message and the new test comment say, and the MLX builder now looks up constants in a list. Detailed reviewPR Review: #24481SummaryThe PR fixes three MLX failures in LeRobot policies:
Nothing blocks merging. The two findings below are minor, and both are in code changed since the last round. Code Quality
Performance
RecommendationApprove. Nothing blocks. Each fix is a few lines. I did not run any tests, because this job can't run Python. Everything above comes from reading the code. |
e6ca3a1 to
dc7d498
Compare
|
Thanks, both fixed. Graph order. Constant lookup. |
|
@claude /pr-review |
|
Claude finished @shoumikhin's task in 24m 1s —— View job Summary: Nothing blocks merging. There are two minor findings, both in the shared-constant copy: the one-pass copy is still quadratic in the number of constants because each one rebuilds the signature lookups, and Detailed reviewPR Review: #24481SummaryThe PR makes Nothing blocks merging. Both findings are minor. Backward Compatibility
Performance
RecommendationApprove. Nothing blocks. Each fix is a line or two. I did not run any tests, because this job can't run Python. Everything above comes from reading the code. |
dc7d498 to
2450a5a
Compare
|
Thanks, both fixed in the copy commit. One tag string. Signature lookups. The candidate check now reads |
|
@claude /pr-review |
|
Claude finished @shoumikhin's task in 24m 43s —— View job Summary: Nothing blocks merging. There are two minor findings: until #24404 lands, an MLX delegate that reads a constant through a list still gets two placeholders with the same target, which is the case the builder branch this PR removes used to handle, and the new mask helper finds and reads the mask in a roundabout way. Detailed reviewPR Review: #24481SummaryThe PR makes four changes:
Nothing blocks merging. Both findings are minor and in the MLX code. Code Quality
RecommendationApprove. Nothing blocks. The builder branch only matters if this PR merges before #24404. I did not run any tests, because this job can't run Python. Everything above comes from reading the code. |
|
Thanks. Both points: Constant read through a list. I built your case: a CPU op reads Mask lookup. Agreed. The cleanup (an explicit |
metascroy
left a comment
metascroy
left a comment
There was a problem hiding this comment.
Looks good overall. I checked that dialect follows the verifier, so delegate submodules, fake programs, and the programs backends get in preprocess all still report EDGE, and the in-tree constant_prop_pass callers keep getting the Edge clone. The grouping and the one-pass copy also look correct. A few things to address inline.
| """ | ||
| A pass to duplicate the attributes/constants node (the candidate_node) in the graph. Mostly used for duplicating light-weight data. | ||
| If the data is too large, try tag it with "no_copy" to prevent high memory usage and make it as part of the output. | ||
| Users with the same delegation tag, or with none, share one placeholder. |
There was a problem hiding this comment.
This changes the behavior of a public function: if no user has a delegation tag, duplicate_constant_node now makes no copies and returns set(), where before every user but one got a copy. In-tree only tests call it directly, but it lives in canonical_partitioners, so out-of-tree partitioners may call it. Could you note the change in the PR description, and maybe say in the docstring that the users must be tagged before calling it?
There was a problem hiding this comment.
Agreed, this is a behavior change. I checked whether to keep the old behavior for untagged users. With the code on main, an untagged lifted tensor constant gets a copy, and then ep.module() fails with "c is not a buffer", the same error this PR fixes. It would also bring back a copy for every CPU user but one. So I kept the new behavior. The docstring now says to tag the users before calling it, and the PR description notes the change.
|
|
||
| if len(non_none) == 1: | ||
| axis, idx = non_none[0] | ||
| idx_meta = n.args[1][axis].meta.get("val") |
There was a problem hiding this comment.
After _constant_mask_as_positions replaces the mask, this line still reads the metadata of the original boolean mask (n.args[1][axis]), not of the positions slot. It is only right because a mask used as the single index is 1-D, like its positions. Could you base the choice between TakeAlongAxisNode and TakeNode on the index actually used, or add an assertion, so a later change to the mask path doesn't silently pick the wrong node?
There was a problem hiding this comment.
Done. Each index now carries its own metadata, and a replaced mask carries its positions. So the choice between TakeAlongAxisNode and TakeNode now reads the index that is actually used. The result is the same today, for the reason you gave.
| target_to_constant_slot[target] = self.make_or_get_slot( | ||
| node, id_space=IdSpace.Constant | ||
| ) | ||
| self.make_or_get_slot(node, id_space=IdSpace.Constant) |
There was a problem hiding this comment.
Removing the target_to_constant_slot workaround is correct with the grouping. But if anything ever gives one delegate two placeholders for the same target again (a direct duplicate_constant_node call, or a future pass), the result is now a second constant slot, which stores the data twice in the .pte without any warning, instead of sharing one slot. Could you add a cheap check here, for example an assertion that each constant target gets one slot, so that would fail loudly?
There was a problem hiding this comment.
Done. _make_io_slots now raises a ValueError when a second placeholder names a constant target that already has a slot. One small correction: without the check, the data would not be stored twice. Constant data is named by target, so one of the two slots would get no data, and the runtime would reject the program when it loads. The new test test_second_placeholder_for_one_constant_raises builds such a program, and it fails when the check is removed.
When ops in different partitions, or in a partition and on the CPU, read the same constant, duplicate_constant_node gives every user but one its own copy of the placeholder, all with the original target. Two users in one partition then bring two placeholders of it into that delegate, and two users left on the CPU keep two in the owning program. A program that lists the same lifted tensor constant twice fails to unlift: ExportedProgram.module() raises "`<name>` is not a buffer". For a delegate, to_executorch raises it too, because it runs the delegate's program when it retraces the delegate call. For example, a constant read by an add and a mul in one delegate and by a sin on the CPU fails this way. So does a second partitioner that takes the users left on the CPU into one delegate. Users with the same delegation tag, or with none, now share one placeholder, and the first group keeps the original. So copies are made only for users in different programs. On a graph without tags none are made, so a caller has to tag the users first. The docstring says so, and the tests that call duplicate_constant_node directly now tag their users. The MLX builder gave placeholders that name the same target one constant slot, to cope with the copy per user. Each delegate now gets one placeholder per constant, so that code is removed, and a second placeholder of the same constant now raises a ValueError. Without it, one of the two slots would get no data, because the data is named by target, and the program would fail only when it loads. The new test lowers a constant (a lifted tensor constant, a buffer, and a parameter) that two delegates each read twice and that ops left in the owning program also read, through two to_backend calls, once with the ops left in the owning program reading it first and once last. It also lowers both delegates through one to_backend call. A direct call on untagged users makes no copy. It checks that each program takes the constant once and that to_executorch passes. An MLX test checks that the builder rejects a second placeholder of one constant.
The MLX partitioner takes aten.index.Tensor whatever the dtype of its
indices. The MLX runtime gathers only with integer indices, so x[mask]
with a boolean mask lowers, then fails when it runs:
MLX execute failed: [Gather] Boolean indices not supported
The mask is often a constant of the program, for example one that
constant_prop_pass folded before lowering. Its values are known at
lowering time, so the index handler now replaces it with its integer
positions, one index per mask dimension, and the gather stays in the
delegate.
Any other boolean index is refused, so the partitioner leaves it on the
CPU:
- a mask computed in the graph, such as x[x > 0]
- a mask that is a method input, or a buffer the program mutates
- a mask that is a parameter or a persistent buffer, because the
partitioner sees those as fake tensors and cannot read the mask
- a constant mask used together with other indices
- an all-false mask, because the result is empty and the MLX runtime
cannot write an empty output
The new tests run constant masks (1-D, 2-D, 3-D, a non-persistent
buffer, and a mask at axis 1), which must stay in one delegate. They
also run a mask computed in the graph, a mask passed as an input, a
constant mask used with another index, a persistent buffer, a buffer
the program mutates, and an all-false mask, none of which may lower the
index. A direct check on the fake program the partitioner sees confirms
that the handler itself refuses the persistent buffer and the all-false
mask, rather than failing on them.
To give each delegate its own placeholder for a constant it shares with another delegate, to_backend called duplicate_constant_node once per shared constant, for each partition tag. Every call recompiled the graph module and validated the whole program, so the cost grew with the number of shared constants times the size of the graph. A model lowered to many delegates, with thousands of folded constants, spent about an hour here. Now to_backend collects the shared constants of all tags first and copies them in one pass with the new duplicate_constant_nodes, which recompiles and validates once. _maybe_duplicate_constant_nodes takes all the tags, and raises a TypeError for a single tag string, which would otherwise be read as one tag per character. The pass reads the kind of each constant from its input spec, instead of through get_param, get_buffer and get_lifted_tensor_constant, which rebuild the signature lookups on every call. So a buffer whose data is not in the state dict, which the verifier allows, no longer raises a KeyError here; its copy needs no data. Each shared constant still gets one copy per other program. The group of its first user in graph order keeps the original. Before, when one op read two shared constants, which delegate kept the original depended on the tag order, on the iteration order of a set of names, and on the order in which earlier copies had rewired that op, so it could change from run to run.
2450a5a to
5e7f110
Compare
|
@claude /pr-review |
PR Review in progress
|
Problem
Some LeRobot policies fail with Apple's MLX backend:
SpecViolationError: Operator '<EdgeOpOverload: aten.clone.default> ...' is not an allowed operator type.to_executorchwithAttributeError: `_prop_tensor_constant551` is not a buffer.[Gather] Boolean indices not supported.to_backendcopying shared constants.The clone error.
constant_prop_passputs a clone in front of a folded program output. The clone was always an Edge op, even on an ATen program, and the MLX partitioner's ATen check rejects it.The "not a buffer" error. When ops in a delegate and ops outside it read the same constant,
duplicate_constant_nodegives every user but one its own copy of the placeholder. A program that lists a lifted constant twice cannot be unlifted, soto_executorchfails.Fix
constant_prop_passchecks the program's dialect. An ATen or training program now getstorch.ops.aten.clone.default, whichto_edgeconverts later.duplicate_constant_nodegroups the users of a constant by delegation tag. Users with the same tag, or with none, share one placeholder. Behavior change: on a graph without tags it now makes no copies, so callers must tag users first. The MLX builder's workaround for several placeholders of one constant is replaced by an error.to_backendcopies the shared constants of all delegates in one pass, with one recompile and validate. The group of a constant's first reader in graph order keeps the original, so the result no longer changes between runs.Testing
to_backendtests have no new failures.to_executorch.#24404 also edits
duplicate_constant_node; whichever merges second needs a small rebase.