Spell the parameter transform as a linear, not an einsum (#1641)

Summary:
Pull Request resolved: #1641

`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

fbshipit-source-id: 969300cb28ea54f588c2a2732dfd0b9482ebe147

Author: billmguo

pymomentum/torch/character.py

@@ -749,11 +749,11 @@ def __init__(
     def forward(self, model_parameters: torch.Tensor) -> torch.Tensor:
         if not hasattr(self, "parameter_transform"):
             raise RuntimeError("Character has no parameter transform")
-        return torch.einsum(
-            "dn,...n->...d",
-            self.parameter_transform,
-            model_parameters,
-        )
+        # Equivalent to einsum("dn,...n->...d", parameter_transform, model_parameters).
+        # Spelled as a linear so torch.export does not decompose it into a pair of
+        # permutes over the [d, n] buffer, which on an exported graph are re-executed
+        # every frame rather than folded away.
+        return torch.nn.functional.linear(model_parameters, self.parameter_transform)
 
 
 class InverseParameterTransform(torch.nn.Module):