[john-rocky]

Summary

Fixes #22078.

The XNNPACK partitioner configs require a static weight (is_param_node in partition/config/gemm_configs.py), so a convolution or a linear whose weight is computed from parameters, which is anything under torch.nn.utils.parametrize such as weight_norm and spectral_norm, is declined and left to the portable kernels together with the weight computation. Nothing warns: WhyNoPartition logs at DEBUG, and the model simply runs slow.

This overrides Partitioner.transform_for_pre_decomposition in XnnpackPartitioner to fold those weights before decomposition, so the fold is XNNPACK-scoped. The one change in to_edge_transform_and_lower is a mark on the constants that more than one method reads. Its signature is unchanged. The hook seeds from the weight and bias positions of the ops the GEMM configs partition (linear, the conv family, addmm, mm, matmul with a 2-d weight), walks back through the constant-only producers, and hands that node set to constant_prop_pass as an allowlist. The walk keeps an explicit stack, so a weight computed through hundreds of producers folds instead of reaching the recursion limit. A parameter-only subgraph anywhere else unlocks no delegation and stays an op; a program in which no weight will fold leaves the hook untouched, without the retrace.

Seven things the hook has to get right at this stage of the pipeline:

• The program is not functionalized yet. A KV-cache update is still an in-place copy_, index_put_ or custom op on the buffer, and the graph signature lists no mutated buffers. Folding a view of a buffer that is written to leaves the write on a constant, and run_decompositions fails on the aliasing (expected compiled_fn to be GraphModule), which is what the Voxtral realtime job hit on an earlier revision. The hook functionalizes first with run_decompositions({}), the call to_edge_transform_and_lower makes right after it anyway, and only when the constant-only walk on the program as given finds a weight that will fold; a quantized model, whose weights arrive through a dequantize, leaves the hook as it came. The one shape this misses is a weight built with an in-place op, impure before functionalization: it stays an op. Cost on a 12-layer, 134M-parameter decoder with 24 cache writes: 0.6 s.
• The skip set. The factory ops that decompose to aten.full (full, new_full, ones, new_ones, zeros, new_zeros) and to aten.full_like (full_like, ones_like, zeros_like), so a scalar fill does not become a stored tensor, for the same reason the pass skips full at the edge level. And every quantized_decomposed / torchao op in the graph, so the Q/DQ chain convert_pt2e or quantize_ leaves on a weight stays in place.
• Buffers are not constants here. The hook sees one method at a time, and a buffer this method only reads can be written by another method of the same program. constant_prop_pass gains fold_buffers (default True, existing callers unchanged) and the hook passes False, so parameters and lifted tensor constants seed the fold and buffers never do; a reader keeps the shared allocation of a buffer another method writes.
• Training graphs are returned untouched. A joint forward-backward program needs its parameters as inputs for the runtime to hand them to the optimizer, the same check to_executorch makes for do_quant_fusion_and_const_prop.
• A fold may not grow the program. Folds that share a source are decided as a group, and the group is applied only if the folded tensors take no more bytes than the sources the fold erases. A parameter used both inside and outside the fold cannot be erased, so its fold would only add a copy: a tied embedding read by one lookup and one transposed matmul stays as it is (.pte unchanged at 133,904 bytes; an earlier revision doubled it).
• A constant another method reads stays out of the fold. Before the pre-decomposition transforms run, to_edge_transform_and_lower marks each constant placeholder that more than one method reads with meta["shared_across_methods"] = True, as documented in Partitioner.transform_for_pre_decomposition. Placeholders are matched by fully qualified name. One with no users is not a read, since a non-strict export lifts every parameter of the module into every method. The hook leaves marked parameters and lifted constants out of the fold, as it does buffers: the other method keeps the original, so a fold would store a second copy. An Embedding(512, 128) read by a lookup in one method and a transposed matmul in another keeps the .pte at 264,720 bytes (an earlier revision stored 526,976). The cost: two methods that fold the same expression over one shared parameter now both keep the op.
• A multi-output op folds with all of its outputs. A packed qkv split or a gate/up chunk in a weight expression brings every getitem with it, so the outputs are materialized one tensor each and the op never is; the byte rule counts tensor leaves, and a scalar that a consumer outside the fold also needs is accounted as the op it stays.

Changes in the pass itself:

• nodes_to_fold, an optional allowlist; fold_buffers, as above.
• register_like_source (default False, so quant_fusion_pass, Qualcomm's convert_mha_to_sha and Samsung's ConstantPropPass see the same names, kind, meta and placeholder order as before). With it, a folded value is a parameter if any of its sources is a parameter, else a buffer if any is one, else a lifted constant. It keeps its sources' custom meta, because run_decompositions carries meta back for parameters and buffers only and a fold lifted as a constant lost its external-file tag. Its name derives from the first source of the fold's kind and from the whole sha256 digest of the expression, w_prop_ followed by 64 hex characters for a value computed from the parameter w, so two methods folding the same expression share one entry in the external constant map and two expressions over one parameter get two; a name from the tensor contents would change with the weights and break loading a .ptd exported from other weights. The expression names each op fully qualified (aten.t.default, _operator.getitem) and describes each producer once, so the string is linear in the expression.
• The pass folds a node only when the placeholders it is computed from agree on their custom meta, where the external-file tag lives. This holds in both registration modes. A merged adapter linear whose base weight is tagged for one file and whose adapter factors for another is therefore not folded whole: the sum stays an op, the base keeps its file, and the adapter product, whose sources agree, folds into a value tagged like them. The hook applies the same rule when it picks the weights to fold.
• Folded placeholders go after the last placeholder whose value lives in the same store, the values are registered once every placeholder exists, in graph order, and an assertion checks that the two orders agree, so a caller that binds the graph module by position (state dict values, constants, inputs) gets the right tensor in every slot.
• The fixup that returns a constant output through a clone builds its constant set with the caller's fold_buffers, so a read-only buffer output is left alone when buffers are not folded.

And one check in the emitter: _save_to_external_constant_map raises when a name already in the map points at different data. The map is shared by the methods of a program and deduplicated by content, so that case is two methods writing different tensors under one name, and the last write was loaded for both.

Earlier revisions: a constant_prop flag on to_edge_transform_and_lower (reworked after review), then the hook without the functionalize step (broke the Voxtral export), then a fold of every parameter-only subgraph (narrowed after review). The impure-op fix the first review found in the pass itself landed as #22418.

Measurements

wav2vec2-large's positional convolution, weight_norm(Conv1d(1024, 1024, 128, padding=64, groups=16)) at sequence length 499, module alone, synthetic weights, macOS arm64, torch 2.13.0:

	delegates	ops left outside	median latency	max abs diff vs eager
main	2	sum, pow, convolution	3912.1 ms	1.6e-05
this PR	1	none	4.7 ms	1.2e-05

.pte size is the same (33.6 MB) in both cases: the folded weight replaces weight_v, it does not join it.

Quantized paths checked with and without the hook, Q/DQ op set and .pte size identical on all seven: legacy 8da4w, quantize_ 8da4w / 8da8w / int4 weight-only, PT2E static / dynamic / qc4w.

Voxtral realtime, CI export path (streaming, 8da4w linears, 8w embedding, XNNPACK) on a 48x scaled-down random-weight instance of the same graph: the earlier revision fails with the CI assertion, this revision lowers, with the same delegate count per method as main (21 / 13 / 0).

Test plan

backends/xnnpack/test/test_xnnpack_partitioner.py:

• test_parametrized_weight_is_folded_before_partitioning: a weight_norm Conv1d lowers to a single delegate call with nothing else at the top level, and the runtime output matches eager.
• test_pre_decomposition_folding_folds_gemm_weights_only: next to a weight_norm conv, an arange mask, an expand and a parameter-derived output stay ops; only the conv weight is folded, the runtime matches eager, and a program with no computed GEMM weight comes back as the same object.
• test_pre_decomposition_folding_does_not_duplicate_shared_weights: a tied embedding and a two-use weight produce the same .pte bytes with the hook as without it.
• test_pre_decomposition_folding_keeps_parameters_shared_across_methods: an Embedding(512, 128) read by a lookup in one method and a transposed matmul in another gives a .pte of the same size as without the hook, the matmul method keeps embedding.weight, and both methods match eager through the runtime. A weight_norm conv weight that only one method reads still folds while the other method reads the conv's bias.
• test_pre_decomposition_folding_handles_multi_output_ops: a packed weight through split, chunk, unbind and max(dim) with one output used outside a GEMM folds with every output, the op is gone, and the runtime matches eager.
• test_pre_decomposition_folding_handles_a_deep_weight_chain: a linear weight built from its parameter by 600 additions folds, the program lowers to a single delegate call, and the runtime matches eager.
• test_pre_decomposition_folding_retraces_only_with_a_fold: a quantized linear, a linear whose weight is an input and a matmul of two activations come back as the same object with run_decompositions never called; a weight_norm conv folds.
• test_pre_decomposition_folding_keeps_quantization_primitives: on a convert_pt2e graph and on a groupwise int4 weight, the set of quantized_decomposed nodes is identical before and after the hook.
• test_pre_decomposition_folding_skips_factory_ops: every op in the skip set is checked to decompose to full / full_like and to stay an op at a linear's weight position.
• test_pre_decomposition_folding_keeps_mutated_buffer: a buffer read and written in place stays a mutated buffer through the hook and through lowering while the weight is folded, and the runtime carries its state across calls like eager.
• test_pre_decomposition_folding_keeps_buffers_shared_across_methods: a buffer one method writes and another reads as a matmul weight comes back from the hook as the same object, stays in the reader's signature, and with share_mutable_buffers=True both methods carry its mem_id=2 allocation.
• test_pre_decomposition_folding_skips_training_graphs: an _export_forward_backward program keeps its parameter and emits __et_training_parameters_index_forward.
• test_pre_decomposition_folding_folds_through_a_scalar_item: a linear weight scaled by scale.item() folds, the item goes with it, and the runtime matches eager.
• test_pre_decomposition_folding_keeps_a_scalar_item_used_outside_the_fold: a float used in the weight expression and outside it: the weight folds, the item and its parameter stay, and the runtime matches eager.
• test_pre_decomposition_folding_handles_view_of_parameter: x @ self.w.t() lowers to a single delegate and matches eager through the runtime.
• test_pre_decomposition_folding_keeps_external_weight_tags: a foundation and a lora weight tagged with delegate_external_constants_pass_unlifted and folded through t() land in foundation.ptd and lora.ptd, with nothing in the .pte.
• test_pre_decomposition_folding_keeps_adapter_and_base_weights_apart: a linear whose weight is a base tagged for foundation.ptd plus an adapter product tagged for lora.ptd comes back from the hook as the same object. foundation.ptd holds the base weight, lora.ptd both adapter factors, and the .pte and both files are the same as without the hook.
• test_pre_decomposition_folding_names_folds_after_their_expression: three methods folding one parameter, two through t() and one through a scaled t(), write two entries into the external constant map, and each method returns its own value through the .ptd.

exir/tests/test_constant_prop_pass.py, new: exir/tests/test_passes.py is ignored by pytest.ini and pytest-windows.ini since #14037, so the tests for the pass options live in a file the exir/ test path collects, with a Buck target. test_constant_prop_pass_folds_through_a_scalar (b * w.max().item() at the edge level comes out as one add against one constant), test_constant_prop_pass_keeps_a_scalar_op_used_outside_the_fold, test_constant_prop_pass_fold_buffers_false, test_constant_prop_pass_nodes_to_fold, test_constant_prop_pass_registers_fold_like_its_source, test_constant_prop_pass_registers_buffer_fold_as_buffer, test_constant_prop_pass_names_folds_after_their_expression, test_constant_prop_pass_describes_each_producer_once (sixteen levels of squaring and a chain of 600 additions; each op is written with its namespace, and the fold's name carries the whole digest), test_constant_prop_pass_registers_folds_in_graph_order (three folds off one parameter, the graph module bound by position, in both registration modes), test_constant_prop_pass_keeps_sources_tagged_for_different_files (a merged adapter weight tagged for two files: the sum stays an op and only the adapter product folds, tagged lora.ptd; with one tag on all three the whole weight folds), test_constant_prop_pass_clones_a_constant_output_in_the_dialect_of_the_graph (aten.clone before to_edge, the edge clone after), test_constant_prop_pass_fold_buffers_false_leaves_a_buffer_output_alone. The existing tests in test_passes.py keep their _prop_tensor_constant0 asserts.

exir/emit/test/test_emit.py: test_constant_tagged_tensor_same_name_different_data.

exir/tests/test_partitioner_pre_decomposition.py: test_constants_read_by_several_methods_are_marked_for_the_transforms (the parameter both methods read is marked in each, the ones only one method reads are not, and a single-method program marks nothing).

Also run locally with the hook in place: backends/xnnpack/test/ops/test_conv1d.py, backends/xnnpack/test/ops/test_linear.py, backends/xnnpack/test/quantizer/test_pt2e_quantization.py, exir/program/test/test_program.py, exir/emit/test/test_emit.py, exir/tests/test_partitioner_pre_decomposition.py, exir/tests/test_passes.py -k constant_prop.

cc @GregoryComer @digantdesai @cbilgin @JakeStevens @nil-is-all

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/22391

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

⚠️ 20 Awaiting Approval

As of commit e212176 with merge base 75af78c ():

AWAITING APPROVAL - The following workflows need approval before CI can run:

• Build Aarch64 Linux Native Wheels (gh)
• Build documentation (gh)
• Build Linux Native Wheels (gh)
• Build Presets (gh)
• Cadence Build & Test (gh)
• Lint (gh)
• periodic (gh)
• pull (gh)
• Test ARM Backend (gh)
• Test CoreML Backend (gh)
• Test Cortex-M Backend (gh)
• Test CPU Backend (gh)
• Test Native Engines (gh)
• Test NXP Backend (gh)
• Test QNN Backend (gh)
• Test QNN Windows MSVC build (gh)
• Test Vulkan Backend (gh)
• Test Vulkan Backend (specialized runners) (gh)
• Test WebGPU Backend (gh)
• Test XNNPACK Backend (gh)

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[JakeStevens]

a few thoughts:

(1) seems there is a bug in the underlying pass for non-determinant ops:

import torch
from executorch.exir import to_edge_transform_and_lower


class RandomAdd(torch.nn.Module):
    def forward(self, x):
      return x + torch.rand(4)


x = torch.zeros(4)
for fold in (False, True):
      ep = torch.export.export(RandomAdd(), (x,), strict=True)
      edge = to_edge_transform_and_lower(ep, constant_prop=fold)
      model = edge.exported_program().module()

      first = model(x)
      second = model(x)

      print(f"{fold=}, outputs_equal={torch.equal(first, second)}")
      print(first)
      print(second)

(2) I don't really like adding a flag to to_edge_transform_and_lower for a pass. maybe better to enable it only for xnnpack in the pre annotation transform. else there may be a better place to hook this in. thoughts maybe @JacobSzwejbka ?

[JakeStevens]

some feedback:

1. The new test imports XNNPACK quantizer and TorchAO, but test_xnnpack_partitioner does not depend on :xnnpack_quantizer or //pytorch/ao:torchao. Add both to its BUCK target.

2. Factory skip list is incomplete. aten.new_full, aten.new_ones, and aten.new_zeros also decompose to aten.full. Add these overloads and a regression test

after these fixes please make sure to run linter as that is failing already

[john-rocky]

Thanks for the review, @JakeStevens. All three points are in, plus one more fix the CI run turned up. The branch is rebased on main now that #22418 landed: the hook (7dec35b) and the review fixes (28837bf).

Your three points. test_xnnpack_partitioner now depends on :xnnpack_quantizer and //pytorch/ao:torchao. new_full, new_ones and new_zeros are in the skip set, and test_pre_decomposition_folding_skips_factory_ops walks the whole set: each op is checked to decompose to full or full_like and to survive the hook. One thing that check showed: full_like / ones_like / zeros_like decompose to aten.full_like, which the edge-level default does not skip. I kept them, since the size argument is the same; happy to trim the set to the edge default if you prefer. flake8 and ufmt are clean on the changed files (the B020 was mine).

The Voxtral job was this PR, not the trunk. The ATen program at this hook is not functionalized yet: the KV-cache update is still an in-place op on the buffer (llama.update_cache_with_indices in Voxtral, index_put_ / copy_ elsewhere), and the graph signature lists no mutated buffers. constant_prop_pass reads mutation from the signature, so it took k_cache / v_cache for constants and folded views of them; run_decompositions then failed on the aliasing between a folded view and the buffer the custom op writes, with the expected compiled_fn to be GraphModule assertion. The hook now functionalizes first with run_decompositions({}), the call _gen_edge_manager_for_partitioners makes right after it anyway; on a 12-layer, 134M-parameter decoder with 24 cache writes that costs 0.6 s. test_pre_decomposition_folding_keeps_mutated_buffer covers it through the runtime, with the state carried across calls. I did not pull the gated checkpoint, so I ran the CI export path (streaming, 8da4w, 8w embedding, XNNPACK) on a 48x scaled-down random-weight instance of the same graph: the previous revision fails with the CI assertion, this one lowers with the same delegate count per method as main.

On the hook location, from your first note: transform_for_annotation only runs under prepare_pt2e / prepare_qat_pt2e, and the case in #22078 is fp32, so transform_for_pre_decomposition is the one spot every XNNPACK export passes through. If you or @JacobSzwejbka would rather have an off switch on the partitioner, that is a small addition.

Of the other red jobs on the previous run: the three unittest jobs died in sccache (AWS credentials) while building the tokenizers wheel, lintrunner-mypy was the backends/arm errors #22482 has since fixed on main, and moshi / qnn are the flaky and trunk failures Dr. CI flagged.

Thanks again for the careful look.

[nil-is-all]

Thanks @john-rocky, running CI again
cc @JakeStevens

[john-rocky]

Thanks for re-running CI, @nil-is-all. Two of the three jobs that failed on that run were my new tests, and the fix is pushed (428ffae); the third, test-models-linux (emformer_transcribe, portable), was a broken pipe while pip installed requirements-examples.txt.

Both tests passed a temporary to_executorch().buffer to _load_for_executorch_from_buffer. The module keeps a pointer into that buffer and loads the method on the first call, so by then the bytes were freed and the outcome depended on what had reused the memory: error: 0x23 on macOS, a worker segfault on Linux, green on the re-run. With MallocScribble=1 both tests fail on every run with Program identifier 'UUUU' != expected 'ET12'; with the program manager held for the life of the module they pass, scribbled or not, and the file is 7/7 locally. The mutated-buffer test also no longer compares the first call, since the initial state of a mutated buffer is not serialized, and both runtime checks use an explicit atol now.

The new commit is waiting on the workflow approval again. Thanks!

[JakeStevens]

findings from an agent-assisted review:

1. [P1] Buffers shared across methods can be frozen. Functionalization only identifies mutations within the current method. If one method writes state and another computes state.sum(),
   the reader’s sum becomes constant. I reproduced the reader returning 0 after a write that should produce 4; disabling the hook restores correctness. The serialized reader also loses
   its shared-buffer allocation. Folding needs to account for mutations across methods.

2. [P1] Training exports lose trainable parameters. For an _export_forward_backward graph computing (x @ weight.t()).sum(), the hook folds the transpose and removes weight, even with
   force_non_static_weights_for_f32_linear=True. The resulting program lacks __et_training_parameters_index_forward, which TrainingModule requires. Skip training graphs or exclude
   parameters with gradient outputs; the existing safeguard in to_executorch() is bypassed by this earlier call.

3. [P2] Scalar constant expressions crash export. A model returning x * self.value.item(), where value is a scalar buffer, fails with TypeError: cannot create weak reference to 'float'
   object. The pass attempts to lift the scalar result as a tensor. Without the hook, lowering and serialization succeed. Preserve these scalar operations or fold their results as
   scalar literals.

[john-rocky]

Thanks for the agent-assisted pass, @JakeStevens. All three reproduce on the previous head, and the fix is pushed (4e7e479, rebased on main): buffers and training graphs stay out of the fold, and the pass no longer lifts a scalar result.

1. Buffers shared across methods. The hook sees one method, so it cannot know which buffers another method writes. It now folds parameters and lifted tensor constants only: constant_prop_pass gains fold_buffers (default True, so the edge-level callers are unchanged) and the hook passes False. In your write/read shape the reader keeps state in its signature, and with MemoryPlanningPass(share_mutable_buffers=True) both methods carry the mem_id=2 allocation for it; test_pre_decomposition_folding_keeps_buffers_shared_across_methods asserts both. The cost is that a subgraph fed only by an unmutated buffer is no longer folded. The case in #22078 is parameters (weight_norm), and the pos_conv shape still lowers to a single delegate with the same parity.

2. Training exports. The hook returns a program with a backward_signature untouched, the check to_executorch already makes for do_quant_fusion_and_const_prop. test_pre_decomposition_folding_skips_training_graphs exports (x @ weight.t()).sum() with _export_forward_backward and force_non_static_weights_for_f32_linear=True: weight stays a parameter and __et_training_parameters_index_forward is emitted.

3. Scalar .item(). This one is in the pass. Before decomposition x * value.item() is mul.Tensor(x, item) with the Python float fed straight in, so the pass tried to lift a float; at the edge level to_edge inserts scalar_tensor, which is why the do_quant_fusion_and_const_prop path never hit it. The pass now records tensor results only, so item and its consumer stay. Covered at the pass level (test_constant_prop_pass_keeps_non_tensor_results) and through the hook with a parameter and a buffer scalar (test_pre_decomposition_folding_keeps_scalar_item).

The seven red jobs on the last run were not this branch. Both lora jobs passed every test (Test 4: Success, Multimethod tests passed!) and then failed in cleanup_files removing the read-only HF cache, which d3a225e fixed on main on Sep 10. nxp-mcuxpresso-test runs a script main added on Sep 4, after this branch's base. test-mcu-cortex-m-backend got a 502 from git.gitlab.arm.com, test_pytest_models_tosa was OOM-killed, unittest-nxp-neutron timed out in the Neutron mobilenet example, which does not import the XNNPACK partitioner, and build-qnn-windows-arm64 is red on main. All of them were green on the previous head 28837bf on Sep 8, and the only change since was the test file. The branch is rebased on main now, so the stale-base ones go away.

Thanks again for the careful look. If you would rather have fold_buffers or the scalar fix as a separate PR, the way #22418 went, I can split it out.

[JakeStevens]

thanks for the back and forth. few more things:

1. a simple x @ self.w.t() fails. see inline comment
2. replace_with_constant_node names lifted constants _prop_tensor_constant{N} from a per-method counter, and copies node.meta from the arithmetic node. This breaks lora on external weights (.ptd). instead, name the lifted constant after its source parameter, and carry the source placeholder's custom meta forward

[john-rocky]

Thanks, @JakeStevens. Both reproduce on the previous head, and the fix is pushed (529dcf8, rebased on main).

1. x @ self.w.t(). Your .detach() is in, applied to every leaf of a folded result. What it fixes: before decomposition aten.t is a real view of the parameter, and a view keeps requires_grad even under no_grad. The retrace in to_edge then cloned it outside no_grad into a non-leaf, and the delegate's copy_exported_program_for_preprocess failed on the deepcopy. The edge-level callers never saw it because to_edge turns t into permute_copy, whose result is a fresh leaf. test_pre_decomposition_folding_handles_view_of_parameter lowers that shape to a single delegate and checks the runtime against eager.

2. Names and meta. replace_with_constant_node now registers the folded value like the first placeholder it is computed from: named {source}_prop (w_prop, then w_prop1 for a second fold of the same source), with the source placeholder's custom meta, and of the source's kind. The kind is the part your note did not ask for, and it turned out to be needed at this hook: run_decompositions, which to_edge_transform_and_lower runs right after the hook, writes the collected meta back for parameters and buffers only (_populate_param_buffer_metadata_to_new_gm), so a lifted tensor constant comes out of it without its custom meta. I checked that on the shape delegate_external_constants_pass_unlifted produces: with the tag carried on a lifted constant, both folded weights still ended up in the .pte; registered as parameters they land in foundation.ptd and lora.ptd. test_pre_decomposition_folding_keeps_external_weight_tags asserts the two files hold the two folded weights and the .pte holds none. test_pre_decomposition_folding_names_folds_after_their_source has two methods whose first fold comes from different parameters and checks the external constant map holds both; with the per-method counter both were _prop_tensor_constant0 and the second write replaced the first.

The naming applies to every caller, so the edge-level tests that asserted _prop_tensor_constant0 now assert a_prop and the like, and test_constant_prop_pass_graph_signature sees a PARAMETER spec. The other callers (quant_fusion_pass, Qualcomm's convert_mha_to_sha, Samsung's ConstantPropPass) run at the edge level, where nothing decomposes afterwards. If you would rather keep the edge-level behavior byte-identical, both the name and the kind can be confined to the hook behind a flag; say the word.

Rerun with the new pass: the partitioner file 16/16, test_passes -k constant_prop 19/19, test_conv1d and test_pt2e_quantization 26/26, test_program 25/25, test_linear with the same three int4 failures my environment shows on main, and the Voxtral, quantized-path and weight_norm checks from the earlier rounds unchanged. The four red jobs on the last run (the apple build, the two webgpu jobs, the windows editable unittest) were infrastructure, not this branch.

Thanks again for the careful look.

[nil-is-all]

@john-rocky could you resolve the failing CI test?

[JakeStevens]

Some notes inline.

[JakeStevens]

as we go back and forth here, I am coming to the idea that we should restrict the blast radius a bit..

Right now, the hook folds every parameter-only subgraph in the graph. But the weight_norm goal is narrower: make a computed weight static so a partitioner config will accept it. A parameter-derived model output, a branch through a custom op with no host kernel, an arange mask, an expand -- none of them feed a static weight, none unlock a delegation, and all of them are now folded and executed on the export machine on every XNNPACK export and hitting some issues in my testing like size regressions: 512× on an expand and 234× on an arange.

One potential approach:

GEMMConfig carries weight_idx and bias_idx, and _get_weight_deps already tests is_param_node on the weight. So the hook can seed from the nodes the enabled configs handle, take the weight and bias positions where is_param_node is false, walk back through constant-only producers, and hand that node set to constant_prop_pass as an allowlist alongside custom_skip_targets.

[john-rocky]

Thanks, @JakeStevens. All four reproduce on the previous head, and the fix is pushed (7a70878, rebased on main). The scope is narrowed as you sketched, and the three findings were mine: the insert point, the duplicated source and the shared name all came out of the round-3 registration change.

1. Scope. The hook seeds from the weight and bias positions of the ops the GEMM configs partition (linear, the conv family, addmm, mm, and matmul with a 2-d weight), walks back through the constant-only producers, and hands that set to constant_prop_pass as nodes_to_fold. A program with no computed GEMM weight leaves the hook untouched, without the functionalize retrace. On a toy with a weight_norm conv next to an arange mask, an expand and a parameter-derived output, the previous head folded all four; now only the conv weight is folded, with the same delegate count (test_pre_decomposition_folding_folds_gemm_weights_only).

2. Insert point. Fixed at the root, both ways you named: a folded placeholder goes after the last placeholder whose value lives in the same store (state dict or constants), the values are registered once every placeholder exists, in graph order, and an assertion at the end of the pass checks that the two orders agree. I cannot run the Qualcomm job here (no QNN bindings), so test_constant_prop_pass_registers_folds_in_graph_order binds the graph module by position the same way, with three folds off one parameter: it fails on the previous head and passes now, in both registration modes.

3. Duplicated source. Folds that share a source are decided as a group, and the group is applied only if the folded tensors take no more bytes than the sources the fold erases. Your two shapes, as .pte bytes with the hook against the hook disabled: a tied embedding (1024x32) went 133,904 to 264,576 on the previous head and stays at 133,904 now; a two-use weight went 68,352 to 133,376 before and stays at 68,352 (test_pre_decomposition_folding_does_not_duplicate_shared_weights). I took the byte count over the live-user rule because it still lets a fold through when a small shared scalar feeds it, and it blocks an expand that would only grow the file.

4. Shared name. The name now derives from the source and from the expression: shared.w_prop_60293755 for w.t() and shared.w_prop_0605881e for (w * 14).t(), the same in every method and export. I did not hash the tensor contents, since that name would change with the weights and a .pte could no longer load a .ptd exported from other weights. Your prefill/decode shape, run through the .ptd: on the previous head one method read the other's weight (4 where eager gives 56); now both match (test_pre_decomposition_folding_names_folds_after_their_expression). As a backstop, _save_to_external_constant_map raises when a name already taken points at different data, with an emitter test.

One step further on blast radius. The round-3 naming and kind change applied to every caller, and it should not have: Samsung's ConstantPropPass keys its quantize_attrs carry on the _prop_tensor_constant prefix, which those names did not carry. Both are now behind register_like_source on the pass; the hook opts in, and the edge-level callers (quant_fusion_pass, Qualcomm, Samsung) get the same names, kind, meta and placeholder order as main. The edge-level tests are back to their _prop_tensor_constant0 asserts.

Rerun: test_passes -k prop 30/30, the partitioner file 18/18, test_emit 69 passed with the one failure unchanged from the previous head, test_conv1d, test_pt2e_quantization and test_program green, test_linear with the same three int4 failures my environment shows on main, and the Voxtral, quantized-path and weight_norm checks from the earlier rounds unchanged.

Thanks again for the careful look.

[JakeStevens]

Some notes inline.

These are about lines the change does not touch, so they could not be attached to one.

In exir/passes/constant_prop_pass.py, around line 634:

This helper always builds the edge clone. That was fine while every caller handed the pass an edge program. The new pre-decomposition hook hands it an ATen program, so a method that returns a parameter or a constant buffer directly ends up with an edge operator among plain ATen ops. It survives the default path only because the caller retraces straight afterwards, which is luck rather than design, and a partitioner listed after this one would see it. Please choose the clone from the dialect the graph is actually in, or skip the fixup before decomposition. While you are in there, this helper builds its constant set with the default buffer setting, so a read-only buffer output still gets a clone even when the caller passed the buffer switch off.

[john-rocky]

Thanks, @JakeStevens. Three of the four inline findings reproduce on the previous head, and the naming one is as you describe. The fixes are pushed, one commit per finding (fca6649, ea434b3, fae30cd, f49140e, b406d10). The branch is rebased on main, which also brings in main's torchao pin.

Both points of the note in your review body are as you describe, and b406d10 fixes them. The helper now picks aten.clone when the program's dialect is ATen or training, and the edge clone otherwise. It also builds its constant set with the caller's fold_buffers, so a read-only buffer output is left alone when buffers are not folded (test_constant_prop_pass_clones_a_constant_output_in_the_dialect_of_the_graph, test_constant_prop_pass_fold_buffers_false_leaves_a_buffer_output_alone).

The details of the four inline findings are on the threads. The shared-parameter fix has one cost: two methods that would fold the same expression over one shared parameter, like the prefill and decode case from an earlier round, now both keep the op. If you would rather fold there, I can make the decision across all methods instead.

Rerun on the rebased head: test_constant_prop_pass.py 12/12 collected through the directory, the partitioner file 21/21, test_partitioner_pre_decomposition.py 2/2, test_passes.py -k constant_prop 17/17, test_emit 69 passed with the one failure the previous head had too, and test_program 25/25.

PR comment (JA)

@JakeStevens さん、ありがとうございます。インラインの4件のうち3件は前の head で再現し、命名の件はおっしゃるとおりでした。修正は指摘ごとに1コミットで push しました(fca6649、ea434b320、fae30cd26、f49140e43、b406d1015)。ブランチは main に rebase 済みで、main の torchao の pin もこれで取り込んでいます。

レビュー本文のメモは2点ともおっしゃるとおりで、b406d1015 で直しました。ヘルパーは、プログラムの dialect が ATen か training なら aten.clone を、それ以外なら edge の clone を選ぶようになりました。定数の集合も呼び出し側の fold_buffers で作ります。そのため buffer を fold しないときは、読み取り専用の buffer の出力には手を付けません(test_constant_prop_pass_clones_a_constant_output_in_the_dialect_of_the_graph、test_constant_prop_pass_fold_buffers_false_leaves_a_buffer_output_alone)。

インラインの4件の詳細は各スレッドに書きました。共有 parameter の修正にはコストが1つあります。1つの共有 parameter に同じ式を fold するはずの2つのメソッドは、今は両方とも op を残します。以前のラウンドで出た prefill と decode のケースがこれに当たります。そこも fold したい場合は、判断をすべてのメソッドをまたいで行う形に変えられます。

rebase 後の head での再実行: test_constant_prop_pass.py はディレクトリ経由の collect で 12/12、partitioner のファイルは 21/21、test_partitioner_pre_decomposition.py は 2/2、test_passes.py -k constant_prop は 17/17、test_program は 25/25 です。test_emit は 69 passed で、失敗は前の head にもあった1件だけです。

[nil-is-all]

@john-rocky could you rebase and ping for review?

[john-rocky]

@nil-is-all Rebased onto main (02b9df3, 2026-10-15). The 17 commits replayed without conflicts and the patch is unchanged. The constant_prop_pass, XNNPACK partitioner and program tests pass locally on the rebased head. @JakeStevens could you take another look when you have a moment?

[john-rocky]

@nil-is-all CI on the rebased head (9acd5c3) has finished: 170 checks passed, 30 skipped, none failed. The PR is ready for review.

[john-rocky]

@nil-is-all Rebased onto main (75af78c). Where main now does the same thing, the PR uses main's code. @JakeStevens #24352 describes itself as a smaller alternative to this PR. Should this stay open too? If it stays, let me know when someone can approve the CI runs and I'll push again. New runs here disappear about 48 hours after their last update.