[JakeStevens]

Summary

Fixes #22078

The XNNPACK partitioner requires a GEMM's weight and bias to be static params. A convolution, linear, addmm or mm whose weight is computed from parameters is declined, and it runs on the portable kernels together with the computation of its weight. torch.nn.utils.parametrizations.weight_norm and spectral_norm produce exactly this graph, and they are common in audio models; the issue measured 941x on wav2vec2's positional convolution.

This is a smaller alternative to #22391 and #23154, which fold parameter-only subgraphs in transform_for_pre_decomposition. Folding there rewrites the shared ATen program, before functionalization and before every other partitioner and method sees it. Here the top-level program is not modified. get_static_param_deps lets a GEMM config accept a weight or bias computed with pure ops from parameters, unmutated buffers or constants. It partitions that computation together with the GEMM and tags it, and FoldComputedParamsPass, the first XNNPACK preprocess pass, folds it into a constant inside the delegate. After the fold, the delegate looks exactly as if the weight had been a parameter, and the named data store's content hashes deduplicate identical folds across methods.

The partitioner only accepts a computation that the fold can replace outright without growing the program. Only the folded tensor and its source parameters may be used outside the computation, and only by that GEMM. The folded tensor may be no larger than the parameters it frees, and sources tagged for an external data file are left alone. A computation shared by several GEMMs is therefore not folded, such as the packed bias split across three linears in #23154; delegating it would require merging those GEMMs into one partition, which can create a cycle. Tags are set while a GEMM is being checked, before it is accepted, so the pass folds a tagged node only when one of the GEMMs that tagged it is in the same delegate. As with a buffer used directly as a weight today, a buffer that only another method writes is treated as constant.

AddmmConfig now also rejects addmm with non-default alpha or beta, which the XNNPACK lowering ignored. On main, torch.addmm(b, x, w, alpha=2, beta=3) with parameters w and b is delegated and returns wrong results.

Suggested review order: get_static_param_deps in backends/xnnpack/utils/utils.py, its two call sites in gemm_configs.py, then fold_computed_params.py.

wav2vec2 pos_conv, weight_norm(Conv1d(1024, 1024, 128, padding=64, groups=16), dim=2) on its own at sequence length 499, synthetic weights, macOS arm64, median of 5 runs after 2 warm-up runs:

	delegates	ops left on portable	median latency	max abs diff vs eager
main	2	pow, sum, convolution	3709.8 ms	1.43e-05
this PR	1	none	3.6 ms	1.45e-05

The .pte is 1,280 bytes smaller: one folded weight replaces weight_g and weight_v, and the program has one delegate instead of two.

Test plan

python -m pytest backends/xnnpack/test/passes/test_fold_computed_params.py

The new tests lower a weight_norm Conv1d and a linear with a computed weight and bias, run them through the runtime against eager, and check that nothing is left outside the delegate. They also check what must not fold: a source used elsewhere, a mutated buffer, a low-rank product that would grow the program, a source that stays another input of the GEMM, an intermediate that also feeds the GEMM, and sources tagged for an external file. A rejected convolution whose weight is computed by a delegated matmul, and a weight computed by an earlier partitioner's delegate, still export and match eager. A param that already has the folded weight's name is not replaced, and a scaled addmm stays on portable and matches eager.

python -m pytest backends/xnnpack/test/passes backends/xnnpack/test/ops backends/xnnpack/test/serialization backends/xnnpack/test/quantizer backends/xnnpack/test/test_xnnpack_partitioner.py -o addopts=""

This includes test_linear.py, test_bmm.py and the quantizer tests, which pytest.ini skips. The failures are the same as on main: three qd8 int4 linear tests, and quantizer tests that fail depending on run order. I also compared against main for weight_norm on Linear, Conv2d, depthwise and transposed convolutions, spectral_norm, a dynamic sequence length, and two methods sharing a folded weight. Tied embeddings, packed biases and low-rank weights produce the same .pte as on main.

This PR was authored with Claude Code.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

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