Transforms: reorder a broadcast constant by permuting, not reinterpreting

backends/transforms/remove_permutes_around_elementwise_ops.py

@@ -744,15 +744,44 @@ def permute_subgraph(self, subgraph: Subgraph) -> bool:  # noqa: C901
                     and const_rank < permute_rank
                     and const_node.meta.get("val") is not None
                 ):
+                    # Broadcasting widens the constant to the region's rank
+                    # before the permutation applies.
                     original_shape = list(const_node.meta["val"].shape)
                     padded = [1] * (permute_rank - const_rank) + original_shape
-                    target_shape = [padded[d] for d in node_end_perm]
-                    target_shape = target_shape[permute_rank - const_rank :]
-                    new_node = graph.create_node(
-                        "call_function",
-                        exir_ops.edge.aten.view_copy.default,
-                        args=(const_node, target_shape),
-                    )
+                    target_shape = [padded[dim] for dim in node_end_perm]
+
+                    # Where each non-unit axis ends up. Unit axes carry no
+                    # elements, so only the order of these decides whether the
+                    # permutation rearranges data or merely reshapes.
+                    destinations = [
+                        node_end_perm.index(axis)
+                        for axis, size in enumerate(padded)
+                        if size != 1
+                    ]
+                    if destinations == sorted(destinations):
+                        # Only unit extents moved, so this is a pure reshape and
+                        # a view says it exactly -- and says it for free, since
+                        # view_copy later becomes a memory.view alias.
+                        new_node = graph.create_node(
+                            "call_function",
+                            exir_ops.edge.aten.view_copy.default,
+                            args=(const_node, target_shape),
+                        )
+                    else:
+                        # Reordering a non-unit extent moves data. A view would
+                        # reinterpret the strides and read different elements,
+                        # so widen with a view and permute at full rank.
+                        widened = graph.create_node(
+                            "call_function",
+                            exir_ops.edge.aten.view_copy.default,
+                            args=(const_node, padded),
+                        )
+                        with graph.inserting_after(widened):
+                            new_node = graph.create_node(
+                                "call_function",
+                                exir_ops.edge.aten.permute_copy.default,
+                                args=(widened, node_end_perm),
+                            )
                 else:
                     continue
             user_node.replace_input_with(const_node, new_node)

backends/transforms/test/test_permute_optimization_passes.py

@@ -2063,6 +2063,91 @@ def test_no_permutes_is_noop(self) -> None:
             count_node(result.graph_module, exir_ops.edge.aten.permute_copy.default), 0
         )
 
+    def _assert_region_cancels(
+        self, module: torch.nn.Module, inputs: tuple[torch.Tensor, ...]
+    ) -> None:
+        """The region's boundary permutes go away and the values do not change."""
+        module = module.eval()
+        expected = module(*inputs)
+        with torch.no_grad():
+            exported = torch.export.export(module, inputs)
+            edge = to_edge(
+                exported,
+                compile_config=EdgeCompileConfig(
+                    _check_ir_validity=False, _skip_dim_order=True
+                ),
+            )
+            before = count_node(
+                edge.exported_program().graph_module,
+                exir_ops.edge.aten.permute_copy.default,
+            )
+            transformed = edge.transform([RemovePermutesAroundElementwiseOps()])
+            actual = transformed.exported_program().module()(*inputs)
+
+        after = count_node(
+            transformed.exported_program().graph_module,
+            exir_ops.edge.aten.permute_copy.default,
+        )
+        self.assertLess(after, before, "the boundary permutes should have cancelled")
+        torch.testing.assert_close(actual, expected)
+
+    def test_lower_rank_constant_reorder_preserves_values(self) -> None:
+        """A broadcast constant is widened and permuted, never reinterpreted.
+
+        A view cannot express the reorder -- it reinterprets strides rather than
+        moving elements -- so each rank below the region's needs its own case.
+        """
+
+        class Rank3(torch.nn.Module):
+            def __init__(self) -> None:
+                super().__init__()
+                self.register_buffer(
+                    "bias",
+                    torch.arange(4 * 8 * 8, dtype=torch.float32).reshape(4, 8, 8),
+                )
+
+            def forward(self, x: torch.Tensor) -> torch.Tensor:
+                return (x.permute(0, 3, 1, 2) + self.bias).permute(0, 2, 3, 1)
+
+        class Rank2(torch.nn.Module):
+            def __init__(self) -> None:
+                super().__init__()
+                self.register_buffer(
+                    "bias", torch.arange(8, dtype=torch.float32).reshape(1, 8)
+                )
+
+            def forward(self, x: torch.Tensor) -> torch.Tensor:
+                return (x.permute(2, 0, 1) + self.bias).permute(1, 2, 0)
+
+        class Rank1(torch.nn.Module):
+            def __init__(self) -> None:
+                super().__init__()
+                self.register_buffer("scale", torch.arange(1.0, 9.0))
+
+            def forward(self, x: torch.Tensor) -> torch.Tensor:
+                return (x.permute(0, 3, 1, 2) * self.scale).permute(0, 2, 3, 1)
+
+        class EqualExtents(torch.nn.Module):
+            """Two non-unit axes of the same size swap.
+
+            The extents read the same before and after, so only their order
+            distinguishes a reshape from a reorder.
+            """
+
+            def __init__(self) -> None:
+                super().__init__()
+                self.register_buffer(
+                    "bias", torch.arange(4, dtype=torch.float32).reshape(2, 2)
+                )
+
+            def forward(self, x: torch.Tensor) -> torch.Tensor:
+                return (x.permute(0, 1, 3, 2) + self.bias).permute(0, 1, 3, 2)
+
+        self._assert_region_cancels(Rank3(), (torch.randn(1, 8, 8, 4),))
+        self._assert_region_cancels(Rank2(), (torch.randn(8, 8, 8),))
+        self._assert_region_cancels(Rank1(), (torch.randn(1, 8, 8, 8),))
+        self._assert_region_cancels(EqualExtents(), (torch.randn(1, 4, 2, 2),))
+
 
 class LayoutPermuteVisibilityTest(unittest.TestCase):
     """The data-movement passes must see both permute dialects.