[sohailshk]

ConvertToLinearPass rebuilds linear from an [add]mm by taking the weight from before its transpose. When constant propagation or a transpose removal pass has already folded that transpose into the constant, the [add]mm reads the transposed weight directly and the rebuilt linear transposes it a second time: square weights give wrong values and other shapes make the pass fail. Only convert nodes whose weight still comes through a transpose, and leave the rest as [add]mm.

AI assistance: most of the code, tests and this description were drafted with Claude Code.

Fixes #10499

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/24784

• 📄 Preview Python docs built from this PR

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

❌ 1 New Failure

As of commit 59bdd6d with merge base e0015e1 ():

NEW FAILURE - The following job has failed:

• Lint / lintrunner-mypy (gh)
  >>> Lint for backends/arm/test/modules/test_moe.py:

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

[meta-cla]

Hi @sohailshk!

Thank you for your pull request and welcome to our community.

Action Required

In order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you.

Process

In order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA.

Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with CLA signed. The tagging process may take up to 1 hour after signing. Please give it that time before contacting us about it.

If you have received this in error or have any questions, please contact us at cla@meta.com. Thanks!

[linux-foundation-easycla]

The committers listed above are authorized under a signed CLA.

• ✅ login: sohailshk / name: sohailshk (605171e)

[meta-cla]

Thank you for signing our Contributor License Agreement. We can now accept your code for this (and any) Meta Open Source project. Thanks!

[nil-is-all]

@mcr229 @digantdesai @metascroy @JakeStevens bringing this to your attention

[sohailshk]

The four red QNN jobs fail while downloading the QNN SDK, before anything in this PR runs: build-qnn-windows-* get InvalidAccessKeyId from S3, and test-qnn / test-backend-linux gets curl: (22) ... 403. This PR only touches the XNNPACK convert_to_linear pass and its test, and the XNNPACK jobs are green.

[JakeStevens]

@claude review effort:max

[claude]

Claude finished @JakeStevens's task in 3m 31s —— View job

---

The fix is correct for folded addmm and mm weights, but it also quietly stops bmm from ever converting and leaves a TODO in builder.py that still points at #10499; both are minor and nothing blocks merge.

Detailed review

PR Review: #24784

Summary

This PR stops ConvertToLinearPass from rebuilding linear from an [add]mm whose weight transpose was already folded into the constant. A leftover [add]mm is still delegated through AddmmConfig/MMConfig, whose visitors emit a fully-connected op with XNN_FLAG_TRANSPOSE_WEIGHTS, so nothing is lost on the main path. There are two minor gaps: bmm and a stale TODO. I couldn't run the tests in this environment, so everything below comes from reading the code.

Code Quality

• backends/xnnpack/_passes/convert_to_linear.py:172: the new filter turns off conversion for bmm.
  • aten.bmm.default is still listed in targets (line 40). But is_transpose (backends/transforms/addmm_mm_to_linear.py:117-125) only accepts a 2-D swap: t_copy, or a permute_copy with dims [1, 0] / [-1, -2]. A 3-D permute such as [0, 2, 1] fails that check.
  • Inside a Linear source partition, a bmm weight usually reaches the op through view_copy(expand_copy(permute_copy(w))), which also fails it.
  • So every bmm in a Linear partition is now skipped, where find() used to walk back to the weight. The fallback apply_addmm_mm_to_linear_transform doesn't handle bmm either.
  • That would be a performance regression (the BMM path instead of fully-connected), not a correctness bug.
  • Low confidence: none of the reviewers confirmed that current export actually puts a bmm inside a Linear partition. If it never does, remove bmm from targets so the list is honest. If it does, the check needs to look at the weight node find() resolves to, not only the direct argument.
• extension/llm/export/builder.py:555-557: stale TODO. It still says "ConvertToLinearPass is not a sound pass and must be called before const propagation" and links ConvertToLinearPass is not sound when transposes are elided #10499, which this PR closes. Please update or remove it in this PR.
• convert_to_linear.py:169-170: comment placement (nit). The comment sits between the for and if clauses of the comprehension, which is awkward to read. It would be clearer above the src_node_dict = assignment.

Testing

• backends/xnnpack/test/passes/test_convert_to_linear.py:46-62: the new test stops at RunPasses.
  • It never partitions or lowers the folded graph. The fix relies on that path: the leftover [add]mm has to go through AddmmConfig/MMConfig and op_addmm/op_matrix_multiplication with a folded [in, out] weight.
  • Suggested changes:
    • Add a variant that runs .partition() / .to_executorch() and checks that the addmm gets delegated.
    • Add a model with one folded and one unfolded linear, to show the unfolded one still converts.

Recommendation

Approve. Ideally the bmm question and the stale TODO get fixed first. Neither is a blocker.

[sohailshk]

@JakeStevens thanks for running the review. I addressed each point in 22b3f54:

• bmm. The concern was right. A plain nn.Linear on a 2-D, 3-D or 4-D input decomposes to view_copy + mm/addmm, but a non-contiguous input (for example self.linear(x.transpose(1, 2))) gives a bmm whose weight arrives as view_copy(expand_copy(permute_copy(w))). The direct-argument check rejected that, so the previous head stopped converting it. The check now walks from the weight argument to the node find() resolves to, and converts when a transpose is on that path. bmm stays in targets.
• bmm with a bias. While testing this I found that on main the same linear with bias=True loses its bias: the add comes after a view_copy, find_bias_for_mm only looks for an add right after an mm, and the pass replaces the partition with a linear that has no bias, so the output is wrong. A bmm whose partition has an add is now left as bmm. Finding that bias properly would be a separate change.
• builder.py TODO: removed, since the pass no longer converts a folded transpose.
• Comment placement: moved above the src_node_dict assignment.
• Tests: the folded [add]mm test now partitions, checks for one delegate and no addmm/mm left, serializes and compares outputs. New: one folded and one unfolded linear in the same model (the unfolded one still converts), and the transposed-input bmm case for bias x fold (only unfolded, no bias converts; all four compare outputs).

test_convert_to_linear.py: 6 passed. With main's convert_to_linear.py, 3 of them fail; with the previous head, the bmm test fails. backends/xnnpack/test/passes: 114 passed; ops/test_bmm.py + ops/test_addmm.py: 8 passed. flake8 and ufmt are clean on the changed files. I ran these against the 1.6.0 wheel with this branch's Python files on top, so no native rebuild.

[JakeStevens]

Thanks for the quick turnaround on the earlier feedback. The bmm bias catch is a good find. One more issue before this lands, plus a few small cleanups.

can_convert walks out of the partition (backends/xnnpack/_passes/convert_to_linear.py:89-93)

find() stops at src_partition.input_nodes + src_partition.params, at placeholders and at dequants. The new loop stops only at a non-call_function node or a dequant, so it can follow args[0] past the partition input and return True on a transpose that sits upstream of the linear. create_linear then builds linear(x, find(weight)) from the partition input, which the mm already used untransposed. That transposes it a second time: the same bug as #10499, through a different path.

Example: F.linear(x, (a.t() + 1).t()), after a transpose-removal pass cancels the paired permutes so the mm reads the add directly. The walk goes add → permute(a) and returns True. Constant propagation alone can't trigger this, because a folded weight becomes a placeholder. But the PR description names transpose removal as a case this fixes, and removing the builder.py TODO implies the pass is now sound. The same unbounded walk can also reach a node whose args[0] isn't a Node (e.g. cat) and raise AttributeError.

Suggested fix: stop at the same set of nodes find() stops at, e.g.

boundary = src_partition.input_nodes + src_partition.params
while (
    weight not in boundary
    and weight.op == "call_function"
    and not is_dequant(weight)
):

The linear's own permute (and view(expand(permute(w))) for bmm) is inside the partition, so convertible nodes are still found. A test with a computed weight that has a transpose upstream would cover this.

Smaller items

• convert_to_linear.py:79-83: the docstring's reason for the bmm guard isn't quite right. find_bias_for_mm is never called for bmm, because create_linear only calls it for mm.default. The guard is correct; just reword the reason.
• convert_to_linear.py:182-183: the inline comment above src_node_dict repeats the can_convert docstring and can go.

[JakeStevens]

lintrunner failure seems to be unrelated

[sohailshk]

Thanks, you were right about the walk. I reproduced it with F.linear(x, (a.t() + 1).t()) after FuseCascadedTransposeOrPermuteOps cancels the permute pair at the weight: the mm then reads the add (a partition input), the old loop went on to permute(a) and returned True, and the rebuilt linear transposed the weight a second time. Pushed 59bdd6d:

• can_convert now stops at src_partition.input_nodes + src_partition.params, as you suggested, and also at anything that is not a Node: the boundary alone doesn't cover a cat outside the partition, whose args[0] is a list and raised AttributeError.
• Two regression tests in test_convert_to_linear.py: the transpose-upstream case (the mm stays an mm, outputs match eager) and the cat case. Both fail with the old loop and pass now.
• Reworded the bmm docstring (the bias add would be dropped because find_bias_for_mm isn't used for bmm) and removed the comment above src_node_dict.

The new tests import executorch.backends.transforms.fuse_cascaded_transpose_or_permute_ops, so if an internal target lists this test's deps it needs //executorch/backends/transforms:fuse_cascaded_transpose_or_permute_ops.