to_edge_transform_and_lower does not constant-fold parameter-only subgraphs, so weight_norm'd convolutions never reach the delegate (941x on one module)

Summary

to_edge_transform_and_lower does not constant-fold subgraphs whose inputs are all parameters, so a convolution or a linear whose weight is computed — anything wrapped in torch.nn.utils.parametrize, which is what weight_norm and spectral_norm are — never reaches the delegate. It falls to the portable kernels silently, and the cost is not small.

ExecuTorch already ships the pass that fixes it. exir.passes.constant_prop_pass is described as "constant propagation for Exported Program with lifted parameters", which is exactly this case. The Qualcomm and Samsung backends call it; XNNPACK's own test/ops/test_conv1d.py runs it as a test stage. The generic lowering path does not.

executorch 1.4.0, torch 2.13.0, Python 3.12, macOS arm64, pip wheel.

Measurement

wav2vec2-large's positional convolution, Conv1d(1024, 1024, kernel_size=128, groups=16) under weight_norm, at sequence length 499 — the module on its own, nothing else in the graph:

                      delegated   median latency
as exported              70.0%      3293.7 ms
+ constant_prop_pass    100.0%         3.5 ms

941x, with the output unchanged (max_abs_diff against eager: 4.578e-05 before, 2.003e-05 after).

What stays outside the delegate without the pass is the weight computation and the convolution it feeds:

{'aten.sum.dim_IntList': 1, 'aten.pow.Tensor_Scalar': 1, 'aten.convolution.default': 1}

XNNPACK's partitioner configs require the weight to be a static parameter (is_param_node in partition/config/gemm_configs.py), and a parametrized weight is not one, so ConvolutionConfig declines and takes the convolution with it. The refusal is logged through WhyNoPartition at DEBUG, so at default log level nothing is printed.

Repro

import torch
from transformers import Wav2Vec2ForCTC
from executorch.exir import to_edge_transform_and_lower
from executorch.exir.passes.constant_prop_pass import constant_prop_pass
from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner

model = Wav2Vec2ForCTC.from_pretrained(
    "jonatasgrosman/wav2vec2-large-xlsr-53-japanese", dtype=torch.float32).eval()
pos = model.wav2vec2.encoder.pos_conv_embed
x = (torch.randn(1, 499, 1024),)

for fold in (False, True):
    ep = torch.export.export(pos, x)
    if fold:
        ep = constant_prop_pass(ep)
    edge = to_edge_transform_and_lower(ep, partitioner=[XnnpackPartitioner()])
    # count executorch_call_delegate vs call_function at top level, then time
    # edge.to_executorch() through the runtime

Whole models pay it in one lump: the same .pte for that model goes from 3485.4 ms to 267.8 ms once the parametrization is out of the graph, on a build whose correlation against eager is 1.000000 either way.

Why it is easy to miss

Nothing looks wrong. The export succeeds, the delegation report shows a healthy-looking 64%, correlation against eager is 1.000000, and the model simply runs slow. The three ops left outside are a pow, a sum and one convolution in a list of hundreds. I spent an afternoon ruling out the attention mask, the portable layer norms, XNNPACK's GEMM path and the attention itself before subtracting the measured parts from the total and going looking for what was missing.

Suggestions

1. Run constant_prop_pass inside to_edge_transform_and_lower before partitioning, or expose it as a flag there. It is already in the tree and already used by two backends.
2. Failing that, a warning when a partitioner config declines an op because its weight is not a parameter would turn this into a one-line fix for the user (torch.nn.utils.parametrize.remove_parametrizations).

Happy to send a PR for either.

cc @JacobSzwejbka @angelayi @GregoryComer @digantdesai @cbilgin @JakeStevens