Repository navigation
Fix crashes, wrong outputs and lowering failures in several backends and in exir - #24222
shoumikhin wants to merge 22 commits into
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/24222
Note: Links to docs will display an error until the docs builds have been completed. ⏳ No Failures, 115 PendingAs of commit 3119872 with merge base 0db7a88 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
|
@claude /pr-review |
|
Claude finished @shoumikhin's task in 22m 42s —— View job Approve: nothing blocks. The main open point is that on OpenVINO before 2026.4, the delegate blob now stays in memory next to OpenVINO's own copy for the whole life of the method. Three smaller robustness and test gaps are noted below. PR Review: #24222SummaryThe PR does two things:
Nothing blocks merging. The OpenVINO memory cost on versions before 2026.4 is worth settling, and three smaller points follow. The PyTorch Inductor and OpenVINO sources are not in this checkout, so I mark statements about their internals as low confidence. I ran no tests. Code Quality
Testing
Performance
RecommendationApprove Nothing here blocks the merge. The one point to settle is the OpenVINO memory cost: either free the blob on versions before 2026.4, or document the cost. That can happen here or in a follow-up.
|
35f7db1 to
55b4840
Compare
|
Thanks for the review. Points 1, 2 and 4 are fixed in the commits that own the code. Point 3 has no code change, see below.
|
55b4840 to
52e8dcf
Compare
52e8dcf to
c39d57f
Compare
c39d57f to
d11bd4b
Compare
|
@claude /pr-review |
|
There are no blockers, so I recommend approving. The nine issues are four edge cases that individual fixes miss, two possible gaps I'm not confident about, two fixes with incomplete tests, and a host round trip on every call for strided CUDA outputs. Detailed reviewPR Review: #24222SummaryThis PR makes 19 independent fixes for crashes, wrong outputs and lowering failures in exir (constant propagation, the Edge verifier, the decomposition table,
I could not run tests in this job. I also did not read the torch, coremltools, OpenVINO or QNN internals, so anything I say about their behavior comes from ExecuTorch's own code and comments. Code Quality
Testing
Performance
RecommendationApprove. Nothing here blocks merging. The two low-confidence items (fixes 6 and 19) need the author's judgment. The rest can be follow-ups, but three are small enough to fold in here: the
|
…kends An AOTI backend (CUDA, Metal) checks at lowering time that every kernel Inductor leaves to ATen exists in the backend's shim library, so a missing one fails the export with a clear list. The check watches two code paths that emit a C shim call, but scatter and index_put fallbacks write their call directly. A program that needs one of those kernels lowers without an error, and the process then dies on the first run with an undefined symbol. For example, Inductor leaves an int64 scatter_add_ to aoti_torch_cuda_scatter_reduce_two_out, which the CUDA shim does not have. Watch those two paths too. They do not pass the kernel name to a method the check can see, so read it from the line they write. This checks the name the generated code really calls.
Since OpenVINO 2026.4, ov_core_import_model wraps the blob it is given in an ov::Tensor and the compiled model can keep reading that memory. Earlier versions read it through a stream during the import only. OpenvinoBackend::init frees the blob right after the import. When the program is loaded from a buffer, Free() does nothing, so nothing goes wrong. When it is loaded from a file, the blob is really freed: with a memory-mapped file the first execute crashes with a segmentation fault, even for a two-layer model, and with a file read into memory the model reads freed memory. Free the blob in init only when the OpenVINO runtime is older than 2026.4, as before. With newer versions keep it: ExecuTorch owns it and frees it after destroy(), so it lives exactly as long as the compiled model. The op tests loaded every program from a buffer, so they never saw this. They now load from a file, as most users do.
…gions
torch.export records a `with torch.autocast(...)` block as a
wrap_with_autocast call on a submodule. DecomposeWrapWithAutocast inlines
that submodule into the graph. It had two bugs:
1. It mapped the submodule's placeholders to the operands by name. When
the region reads inputs or parameters, the placeholders are named like
arg2_1, arg0_1, arg1_1, not after the operands, so the pass failed with
a KeyError such as 'arg2_1'. For example:
with torch.autocast("cpu", enabled=False):
y = self.linear(x)
2. It only looked at the top graph. A region nested in another region was
copied out with its get_attr, whose submodule name only means something
inside the outer submodule. In the top module it resolved to the outer
body, so the program silently ran the wrong code.
Map placeholders to operands by position, and inline nested regions
before their parent.
Edge turns `x.detach()` into `aten.detach_copy`. Inductor has no lowering
or C shim for it, so an AOTI backend (CUDA, Metal) cannot lower any graph
that calls detach, for example:
return (x * 2).detach() + 1
fails with "missing fallback kernels: aten::detach_copy".
ReplaceViewCopyWithViewPass already maps copy ops Inductor cannot lower
(slice_copy, select_copy, split_copy) back to views. Map detach_copy to
alias, which is what detach means without autograd. Mapping to detach
itself would still fail Edge verification under the default compile
config, because detach is a non-core view op that the partitioner
preserves.
AOTInductor can return a method output as a view into a larger buffer with strides that do not match a contiguous tensor, for example a slice `y[:, :, :14]` of a [1, 50, 32] result. The CUDA backend copies outputs into the planned ETensor only from contiguous SlimTensors, so execute() failed with "SlimTensor must be contiguous" although lowering accepted the program. The copy already has a slow path that rearranges elements through host memory when the strides differ from the ETensor's. Let it take non-contiguous views too: copy the span the view reaches instead of nbytes(), and check that span against the storage, with checked arithmetic. Wrapping a SlimTensor without a copy still requires a contiguous one. A CUDA graph replay copies each output's raw bytes. For such a strided output, and also for a contiguous output whose ETensor uses another dim order, the replay was silently wrong. A raw copy is known to be right when both are contiguous in the default dim order, so for any other layout warmup now turns graphs off and the capture step refuses to replay, so the method runs without a graph.
torch.export keeps in-place ops in the program it returns. In
t = torch.zeros(6)
t[1:3] = x[:2]
return t * 2
`t` and its slice look constant to constant_prop_pass, so the pass turned
them into constants, while `copy_` still writes into them at run time.
Decomposing the program then fails with "Constant ... is mutated in the
forward method". A buffer written in place has the same problem, and
there the folded reads are silently wrong, because buffers_to_mutate
does not list it until the program is decomposed.
Find the tensors an op writes in place (through any argument, including
out= and lists, as insert_write_back_for_buffers_pass already reads
them), directly or through a view, and do not fold them or their reads.
Programs without in-place ops, such as Edge programs, are unchanged.
The MLX partitioner decides which nodes the delegate can take, and part of that is rejecting nodes whose inputs have a dtype MLX cannot hold, such as float64. The check only looked at arguments that are a single node, so a tensor list argument (the inputs of aten.cat, for example) was never checked. As a result a cat of float64 tensors was put into the delegate. Its float64 inputs became delegate inputs, and lowering then failed with "placeholder has unsupported dtype: torch.float64" instead of leaving the cat on the CPU. Sinusoidal timestep embeddings computed in float64 hit this. Check every node inside each positional argument, so such a cat stays on the CPU like any other float64 op.
When a weight is used by several partitions, to_backend copies its placeholder once per extra user. Every copy keeps the original target, so one MLX subgraph can get several placeholders for the same weight (a layer applied at more than one step, for example). The builder made a separate constant slot for each of them, but the named slots are keyed by target, so only one of those slots got a name. Loading the program then failed with "load_constants: no name found for constant tensor N". Give placeholders that share a target one constant slot. The weight is then stored and loaded once per subgraph.
coremltools has no bfloat16 type. The support query that CoreMLPartitioner
asks coremltools only checks the op name, so an op that reads or produces
a bfloat16 tensor was still delegated, and coremltools 9.0 then raised
"KeyError: torch.bfloat16" during lowering. The whole export failed instead
of leaving that op on the CPU. Models that run part of their forward pass
under bfloat16 autocast hit this.
Override support for nodes with a bfloat16 input or output, the same way
the partitioner already skips other ops coremltools cannot convert.
That rejection is per call, but ops_to_not_decompose keeps an op whole as
soon as one of its calls is supported. With attention run once in float32
and once in bfloat16, the bfloat16 call was kept whole, rejected by the
partitioner, and left with no kernel ("Missing out variants"). Keep an op
whole only when every call of it is supported, as the MLX partitioner
does.
A Core ML tensor has at most 5 dimensions. The support query that CoreMLPartitioner asks coremltools only checks the op name, so an op on a rank 6 tensor (for example a stack of 5-D tensors) was still delegated, and coremltools 9.0 then raised "Core ML only supports tensors with rank <= 5" during lowering. The whole export failed instead of leaving that op on the CPU. coremltools does rewrite layout chains through a higher rank, such as the view, permute, view of a window partition, as long as each step has static shapes and a single user. So override support for a node when the tensors above rank 5 it is connected to reach anything other than such a chain of view, permute and clone ops. Those ops then run on the portable backend, and window partitions stay delegated as before.
A Vulkan program failed to load when a delegated op used a constant tensor with more than 4 dims, for example x * c with a 5-D buffer c: Tensor.h:790: (sizes_.size() <= 4) is false! The Vulkan backend stores tensors with more than 4 dims in buffers, and the partitioner delegates such ops on purpose. 5-D activations already work. But add_prepack_standard_node, which uploads a constant at load time, still passed the tensor's sizes, strides and numel as push constants for buffer storage. Those 4-D push constants assert on a 5-D tensor, and the nchw_to_buffer shader no longer reads them: it gets everything from the BufferMetadata UBO, the same way add_staging_to_tensor_node already works. Drop the unused push constants for buffer storage.
The TOSA backend skipped every top-level placeholder without users. That is fine for constants, but a delegate input can also lose all its users during the Arm pass pipeline. Two examples: - ConvertFullLikeToFullPass turns zeros_like(y), ones_like(y) and full_like(y, v) into full, so y is no longer used. - RemoveNoopPass folds torch.empty(...).copy_(x) into x, so the empty tensor, which stays on the CPU, is no longer used. The runtime still passes every delegate argument, and the VGF backend binds them by position. With the unused input missing from the graph, inputs after it were read from the wrong slots, and outputs were written one slot early, into the last input tensor. The program ran without an error but returned wrong values (all zeros in the copy example). Keep unused placeholders that are delegate inputs. The model converter can still remove dead inputs, which the VGF runtime already handles. Ethos-U is not fixed by this change: Vela drops the unused input as well, and the Ethos-U runtime also binds arguments by position. The new test case is marked as an expected failure there.
NeutronQuantizer did not work on programs from a non-strict torch.export (the default of torch.export.export), because two places that read the source_fn_stack node metadata assumed it is there, and a non-strict export does not record it: 1. get_aten_node_target_partitions skipped every node without it. The same conv, relu and linear model got 10 quantize and dequantize nodes with strict=True and 2 with strict=False, so almost nothing was delegated. The function only uses the metadata to group partitions, and its caller flattens the groups anyway. 2. MoveActivationBeforeConcat, which runs before annotation, raised KeyError: 'source_fn_stack' on a conv, cat, relu pattern. Group by the node's target when the metadata is missing, and copy the metadata to the new activation nodes only when it is there. The new test quantizes a non-strict conv, cat, relu model, which goes through both.
XNNPACKQuantizer annotated ops that read or produce bfloat16 tensors, such
as linears in a model that casts its activations to bfloat16. The quantize,
dequantize and choose_qparams nodes it produced then ran on bfloat16
tensors. XNNPACK does not take those patterns, so they stayed on the
portable quantized kernels, which do not take bfloat16 and abort the
process at run time ("Expected input to be Float tensor received: 15").
Skip nodes that read or produce a bfloat16 tensor, both in the filter that
every annotator applies and when annotations are propagated to ops that
share their input's quantization parameters. Such ops stay unquantized and
run correctly. The cat annotator did not apply that filter at all; it now
does, like the other annotators, so cats now also respect module and
filter scoping. Otherwise float32 and float16 models quantize as before.
to_backend could fail with "Expected output node, got op=call_function" when a runtime check outside a partition reads a value the partition computes, for example torch._check(t.all().item()) where the tensor ops are delegated and the item() and the check are not. After fusing the partition, create_submodule_from_nodes re-sorted the graph with legalize_graph. That sort places a node as soon as its inputs are ready, so the output node could be placed before check nodes that have no users. Later code calls graph.output_node(), which requires the output node to be last. Use stable_topological_sort instead. It keeps the original order wherever the dependencies allow it, so the output node stays last. This changes the order of the remaining top-level nodes of every program that to_backend touches: they now keep their original order instead of the order the old sort produced.
A model that uses a tensor-wide check, for example torch._check_tensor_all or transformers' torch_compilable_check, could not go through to_edge. Two things stopped it: 1. torch.export turns the check into aten._is_all_true, which is not a core ATen op and has no decomposition, so the ATen verifier rejected it. 2. The check then reads the result with item() and compares it through torch.sym_ite. The Edge verifier rejected every Python symbolic function (torch.sym_ite, sym_max, sym_min, sym_not, sym_float, sym_int), on data-dependent values and on shapes alike, even though ExecuTorch lowers each of them to its own executorch_prim op and the runtime runs it. Like other runtime asserts, the check is removed later by to_executorch unless a backend takes it into a delegate; the problem was that the program never got that far. Decompose aten._is_all_true to aten.all in the default decomposition table, and let the Edge verifier accept the Python symbolic functions that ExecuTorch maps to prim ops, the same way it already accepts those prim ops.
A program lowered with the CUDA backend could pass lowering and then kill the process on its first run: undefined symbol: aoti_torch_as_strided Inductor's C++ wrapper calls aoti_torch_as_strided when it allocates a buffer with a larger shape than its logical shape and then gives it the logical shape. ExecuTorch's CUDA shim library did not provide it, so the compiled model could not be loaded. Add aoti_torch_as_strided with the PyTorch shim's signature. It returns a view with the given sizes and strides and the tensor's own storage offset, built on the existing aoti_torch__reinterpret_tensor.
… other Lowering a dynamically quantized pi0 model to XNNPACK failed with "Invalid partition, found dependency cycles". Before merging two partitions, GroupBasedPartitioner checks that no node outside them is both downstream and upstream of the merged partition. That check looks at single nodes. But every other partition also becomes a single node once it is fused. So a path can leave the merged partition, enter another partition at one node, leave it from a different node that the first one does not reach, and come back. The two partitions then depend on each other, and fusing them fails. Before accepting a merge, also follow paths through the other partitions: whenever the merged partition reaches a node of another partition, add everything that partition reaches as a whole, and reject the merge as soon as that reaches something upstream of it. Only partitions that will be returned count, so partitions that are later filtered out (for example ones without enough compute nodes) do not block a merge. Their compute node count now comes from one helper shared with the filter. To stay fast on large graphs, the check only visits the partitions it actually reaches, and caches what each partition reaches; a merge only drops the cached entries that reach the merged nodes. On synthetic graphs of 4k to 16k nodes the partitioner takes about twice as long as before. The merge loop itself is unchanged and still tries every remaining pair, so a graph where the check rejects most merges now keeps many partitions and pays for all their pairs (about 2 s for 1600 partitions on a 4k-node chain of groups); before, such graphs had few partitions only because merges that cannot be fused were accepted. If caller-provided node groups already depend on each other, they are returned as before and fusing them still reports the cycle; merges whose paths reach such groups are refused. Because a merge that would need a later merge to become valid is now rejected, a few graphs end up with a different number of partitions.
The Qualcomm backend has no bfloat16 tensor type (QNN_TENSOR_TYPE_MAP has no entry for it), but nothing checked data types before a node visitor was asked to define the node's tensors. An unquantized bfloat16 node (for example an explicit x.to(torch.bfloat16) region in a float model) therefore crashed lowering with KeyError: torch.bfloat16 in NodeVisitor.get_data_type instead of staying on CPU. Add has_unsupported_dtype, true when the node's output or one of its inputs is an unquantized tensor of a data type the backend cannot represent, and use it in two places: - The partitioner skips such a node, so it falls back to CPU like any other unsupported node. - filter_fn does not keep such a node from decomposition when it is outside a quantized region. Otherwise an op the backend normally lowers whole, such as a bfloat16 linear, would stay on CPU undecomposed and fail with "was not decomposed or delegated". A quantized tensor takes its QNN type from its encoding, so quantized bfloat16 regions are still delegated as before. A bfloat16 graph input to a quantized model is still not supported.
MatmulToBmmPass read the operand shapes with ProxyValue.to_tensor(). When
an operand of the matmul is a constant (a parameter, a buffer or a tensor
attribute), ExportPass passes it to call_operator as a plain torch.Tensor,
so CortexMQuantizer crashed in prepare_pt2e for every such matmul:
AttributeError: 'Tensor' object has no attribute 'to_tensor'
For example softmax(x) @ codebook, where codebook is a tensor attribute.
Read the shapes through .data, which works for both kinds of operand. A
constant operand of a shape the pass rewrites, such as a batched rank-3
codebook used as the right-hand side, now becomes bmm like any other operand
and lowers to cortex_m.quantized_batch_matmul, which already supports a
constant right-hand side. Shapes the pass does not rewrite, such as
rank-3 @ rank-2, no longer crash but stay matmul.
When ExportPass meets a real tensor, for example a get_attr constant, it
makes a fake tensor for the node metadata. Only parameters were faked with
static_shapes=True. Any other tensor that the fake mode had not seen before
got a fresh symbol for every dimension when the fake mode has a ShapeEnv,
which is the case for exported programs.
Such tensors are common: constant folding (for example the Arm quantizer's
ConstantFoldingPass) creates new constant tensors, and the next ExportPass
gives them symbolic sizes like (s45, s96) instead of (9, 2). Later passes
then see symbolic shapes where everything is static. With the LeRobot
diffusion policy, the Arm quantizer then failed in DecomposeVarPass:
AssertionError: a.node.constant should not be None
A real tensor is constant data with known sizes, so fake all of them with
static shapes, as the existing TODO suggested.
aten.slice_scatter is functional: it returns a new tensor and leaves its base alone. The MLX builder lowers it to a SliceUpdateNode whose output slot differs from its base slot, and the runtime handles that case without writing the base. Still, exec_slice_update always fetched the base through tensor_ref(), which is the accessor for tensors an op writes, and which refuses constants. So a slice_scatter whose base is a buffer or another constant failed at run time with "tensor_ref: cannot mutate constant tensor". Read the base with const_tensor_ref(), and use tensor_ref() only for the in-place form, which is the one that writes it.
d11bd4b to
3119872
Compare
|
Please hold off on merging this one for now. I'm adding a last batch of fixes over the next few hours and will comment here when it is final and ready to merge. |
Bug fixes found while exporting models to ExecuTorch backends, one commit per fix. AOTI (AOTInductor) is the compiler behind the CUDA and Metal backends.
Fixes
initno longer frees model data the compiled model still reads.DecomposeWrapWithAutocastmatches submodule inputs by position and inlines nested regions first.detach_copymaps toalias.constant_prop_passno longer folds tensors or buffers that schema-marked in-place writes change, directly or through views.torch.exportprograms.catnow follows module scoping too.to_backendkeeps the output node last when it fuses a partition.to_edgelowerstorch._check_tensor_allandtorch.sym_ite. Partitioned programs also need fix 15.aoti_torch_as_stridedshim.GroupBasedPartitionerno longer merges partitions into mutually dependent ones.MatmulToBmmPasshandles a constant operand.ExportPassgives tensor constants static shapes.For release/1.6: fix 6 needs #24224, fixes 16 and 18 need #24247, fix 19's test needs #23915, and fix 21's test needs #23809.
Testing
Each fix adds or extends a test that fails without it. CUDA and Vulkan ran on an A100 GPU, OpenVINO on an x86 CPU with 2026.4.1 and 2026.3.0, Qualcomm with a QNN SDK build, Core ML and MLX on macOS, Arm with the TOSA reference model and the VGF emulation layer, and NXP with the eIQ Neutron SDK. Metal was not run. GitHub CI does not run the tests for fixes 6 and 11, and skips fix 19's test without the QNN SDK. All three ran by hand.