diff --git a/backends/transforms/remove_permutes_around_elementwise_ops.py b/backends/transforms/remove_permutes_around_elementwise_ops.py index 34068e97ecd..178e6aa3c72 100644 --- a/backends/transforms/remove_permutes_around_elementwise_ops.py +++ b/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) diff --git a/backends/transforms/test/test_permute_optimization_passes.py b/backends/transforms/test/test_permute_optimization_passes.py index ed5578f7fb0..c40061a45db 100644 --- a/backends/transforms/test/test_permute_optimization_passes.py +++ b/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.