[billmguo]

Summary:
ParameterTransform.forward computes
einsum("dn,...n->...d", parameter_transform, model_parameters). With the
transform shaped [out, in], that contraction is exactly x @ W.T -- i.e.
torch.nn.functional.linear(model_parameters, parameter_transform).

The two are numerically equivalent (measured bit-identical on CPU float32 across
unbatched and batched inputs), but they lower very differently under
torch.export. The einsum survives export as a single aten.einsum.default, and
to_edge then decomposes it into unsqueeze -> permute -> bmm -> permute. Both
permutes are over the parameter_transform buffer, which is registered with
requires_grad=False and is identical on every invocation -- so an exported model
re-materializes the entire transform matrix twice per frame. For a character with
889 joint parameters and 204 model parameters that is 181,356 elements copied
twice, on every frame, forever.

Spelling it as a linear removes the output permute and both activation reshapes.
It also lets backends that have a fully-connected primitive (XNNPACK, for
instance) match the op and pre-pack the weight once at initialization rather than
transposing it per invocation, since xnn_define_fully_connected already wants
the [out, in] layout the buffer is stored in.

Measured on an exported production model running on an Android device:

• The block containing this op drops 30.4% (4.50 ms -> 3.13 ms per
  invocation). Within it, permute time goes 1.86 ms -> 0.61 ms and the bmm is
  replaced by an equivalently-priced mm, so the win is entirely the eliminated
  data movement.
• End-to-end on the full exported model, -7.8%, confirmed with a paired
  interleaved A/B in which every measured block favored the change on both the
  median and the minimum.

Eager-mode behavior and results are unchanged; this only affects how the module
lowers when exported.

One consequence worth stating for anyone exporting a model that uses this: the
lowered kernel changes (bmm -> mm / fully-connected). Eager output is
bit-identical, but two different GEMM paths are not bit-identical to each other,
so a model whose output feeds back into its own next-frame input can accumulate a
small difference over a long rollout. Re-validate such a model's output rather
than assuming exact parity with a previously exported artifact.

Reviewed By: nickyhe-gemini, cdtwigg

Differential Revision: D118714733

[meta-codesync]

@billmguo has exported this pull request. If you are a Meta employee, you can view the originating Diff in D118714733.

[meta-codesync]

This pull request has been merged in 8e7451b.