Phonon: run the Mimi decoder's transposed convolutions as convolutions on XNNPACK (#25004)

examples/models/phonon/README.md

@@ -57,7 +57,9 @@ cmake-out/examples/models/phonon/phonon_runner --model_path=... --text="..." --s
 ```
 
 XNNPACK delegates the linear and convolution ops; the flow LM's attention runs
-as the ExecuTorch custom op `llama::sdpa_with_kv_cache`.
+as the ExecuTorch custom op `llama::sdpa_with_kv_cache`. The export also
+rewrites the Mimi decoder's transposed convolutions as convolutions, which
+XNNPACK runs faster.
 
 `--quantize 4bit` or `--quantize 8bit` quantizes the linear layers of the flow
 LM's transformer and flow net with the scheme of the backend; the codec stays in

examples/models/phonon/export_phonon.py

@@ -82,6 +82,7 @@ def source_transforms_for(backend: str, webgpu_compat_ops: bool):
             raise ValueError("--webgpu-compat-ops only applies to --backend webgpu")
         return [st.use_webgpu_compat_ops]
     return {
+        "xnnpack": [st.depthwise_conv_transpose_as_mul, st.conv_transpose1d_as_conv1d],
         "mlx": [st.decompose_elu],
         "vulkan": [
             st.decompose_elu,

examples/models/phonon/source_transformations.py

@@ -4,7 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-"""Function-preserving module swaps for backends that lack an op.
+"""Function-preserving module swaps for backends that lack an op or run it
+slowly.
 
 The model in model.py uses plain PyTorch ops. Each swap below computes the
 same function from more widely supported ops; export_phonon.py selects them
@@ -13,10 +14,15 @@
   decompose_elu           MLX and Vulkan have no aten.elu kernel; without it
                           the Mimi decoder splits at every ELU.
   conv_transpose1d_as_2d  Vulkan runs transposed convolution on 4-D inputs only.
+  conv_transpose1d_as_conv1d
+                          XNNPACK's deconvolution is slow; this runs the Mimi
+                          decoder's transposed convolutions (kernel =
+                          2 * stride) as a convolution and an interleave.
   depthwise_conv_transpose_as_mul
                           Vulkan computes the Mimi upsampler (depthwise,
                           kernel = 2 * stride) incorrectly; this runs it as a
-                          multiply and an overlap-add.
+                          multiply and an overlap-add. XNNPACK also runs it
+                          faster this way.
   use_webgpu_compat_ops   WebGPU runtimes built before these runtime fixes
                           (off by default; --webgpu-compat-ops):
                             * native_layer_norm resize hook (the flow LM's
@@ -93,6 +99,36 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return y.squeeze(2)
 
 
+class ConvTranspose1dAsConv1d(nn.Module):
+    """ConvTranspose1d with kernel_size == 2 * stride as a kernel-2 Conv1d
+    with stride * out_channels outputs (one per output phase), interleaved
+    into the output steps."""
+
+    def __init__(self, convtr: nn.ConvTranspose1d):
+        super().__init__()
+        (self.stride,) = convtr.stride
+        s = self.stride
+        w = convtr.weight.detach()  # [cin, cout, 2s]
+        cin, cout, _ = w.shape
+        # Phase j of output step p reads input p - 1 through tap j + s and
+        # input p through tap j.
+        taps = torch.stack([w[..., s:], w[..., :s]], -1)  # [cin, cout, s, 2]
+        self.weight = nn.Parameter(
+            taps.permute(1, 2, 0, 3).reshape(cout * s, cin, 2).contiguous(),
+            requires_grad=False,
+        )
+        self.bias = convtr.bias
+        self.out_channels = cout
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        z = F.conv1d(x, self.weight, padding=1)  # [b, cout * s, n + 1]
+        b, _, n = z.shape
+        s = self.stride
+        y = z.view(b, self.out_channels, s, n).transpose(2, 3)
+        y = y.reshape(b, self.out_channels, n * s)
+        return y if self.bias is None else y + self.bias.view(1, -1, 1)
+
+
 class ComputeIn(nn.Module):
     """Runs `module` with its weights in `dtype` on a float32 input and
     returns float32. Other attributes (bias, kernel_size, ...) are read from
@@ -171,6 +207,23 @@ def conv_transpose1d_as_2d(programs: dict[str, nn.Module]) -> int:
     return count
 
 
+def conv_transpose1d_as_conv1d(programs: dict[str, nn.Module]) -> int:
+    count = 0
+    for program in programs.values():
+        for module in program.modules():
+            if not isinstance(module, StreamingConvTranspose1d):
+                continue
+            convtr = module.convtr
+            if (
+                isinstance(convtr, nn.ConvTranspose1d)
+                and convtr.groups == 1
+                and convtr.kernel_size[0] == 2 * convtr.stride[0]
+            ):
+                module.convtr = ConvTranspose1dAsConv1d(convtr)
+                count += 1
+    return count
+
+
 def depthwise_conv_transpose_as_mul(programs: dict[str, nn.Module]) -> int:
     count = 0
     for program in programs.values():

examples/models/phonon/tests/test_phonon.py

@@ -43,6 +43,7 @@
 )
 from executorch.examples.models.phonon.source_transformations import (
     compute_weights_in,
+    ConvTranspose1dAsConv1d,
     use_webgpu_compat_ops,
 )
 from executorch.examples.models.phonon.tests.reference import (
@@ -236,19 +237,43 @@ def test_webgpu_compat_ops_preserve_outputs(self):
                 for a, b in zip(outputs, programs[name](*inputs[name])):
                     self.assertLess(max_diff(a, b), 1e-4)
 
-    def test_vulkan_swaps_preserve_outputs(self):
+    def _check_swaps(self, backend: str, counts: list[int]):
         inputs = example_inputs(TINY)
         programs = build_programs(self.model)
         with torch.no_grad():
             before = programs["mimi_decode"](*inputs["mimi_decode"])
-        # 7 ELUs, the depthwise upsampler, 2 decoder transposed convs.
-        counts = [t(programs) for t in source_transforms_for("vulkan", False)]
-        self.assertEqual(counts, [7, 1, 2])
+        self.assertEqual(
+            [t(programs) for t in source_transforms_for(backend, False)], counts
+        )
         with torch.no_grad():
             after = programs["mimi_decode"](*inputs["mimi_decode"])
         for a, b in zip(before, after):
             self.assertLess(max_diff(a, b), 1e-4)
 
+    def test_vulkan_swaps_preserve_outputs(self):
+        # 7 ELUs, the depthwise upsampler, 2 decoder transposed convs.
+        self._check_swaps("vulkan", [7, 1, 2])
+
+    def test_xnnpack_swaps_preserve_outputs(self):
+        # The depthwise upsampler, 2 decoder transposed convs.
+        self._check_swaps("xnnpack", [1, 2])
+
+    def test_conv_transpose_as_conv_matches(self):
+        torch.manual_seed(0)
+        for stride, steps, bias in (
+            (1, 1, True),
+            (4, 16, False),
+            (5, 1, True),
+            (6, 16, True),
+        ):
+            convtr = torch.nn.ConvTranspose1d(
+                6, 3, 2 * stride, stride=stride, bias=bias
+            )
+            x = torch.randn(1, 6, steps)
+            with torch.no_grad():
+                got = ConvTranspose1dAsConv1d(convtr)(x)
+                self.assertLess(max_diff(got, convtr(x)), 1e-5, stride)
+
     def test_compute_weights_in_keeps_float32_io(self):
         inputs = example_inputs(TINY)
         programs = build_programs(self.model)