Arm backend: match operand ranks after view/permute propagation

The view/permute propagation passes can rank-broadcast one operand of a binary elementwise op — a per-channel affine's activation becomes rank 4 while its scale/bias constant stays rank 1 — which TOSA rejects, so Vela/Regor fails to compile the graph. Running MatchArgRanksPass after propagation re-matches operand ranks and keeps the graph TOSA-legal. Adds a regression test that lowers a per-channel-affine module for Ethos-U55.

Authored with assistance from Claude.

Author: rascani

backends/arm/_passes/arm_pass_manager.py

@@ -633,6 +633,9 @@ def _tosa_pipeline(
                 FuseViewCopyTransformPass(),
                 PropagateViewCopyPermuteDownPass(self.compile_spec, exported_program),
                 PropagateViewCopyPermuteUpPass(self.compile_spec, exported_program),
+                # Propagation can leave a binary op with mismatched operand ranks,
+                # which TOSA rejects; re-match ranks before lowering.
+                MatchArgRanksPass(exported_program),
                 RewriteHighRankSingletonPermutePass(),
                 DecomposePermuteForU55Pass(),
                 RewriteSlicePass(),

backends/arm/test/passes/test_rank_match_after_propagation.py

@@ -0,0 +1,51 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+"""Regression test: view/permute propagation can leave a binary op with
+mismatched operand ranks (rank-4 activation vs rank-1 per-channel constant),
+which TOSA rejects. MatchArgRanksPass must run after propagation to keep the
+graph TOSA-legal, so this guards that a per-channel-affine module lowers for U55.
+"""
+
+from typing import Tuple
+
+import torch
+
+from executorch.backends.arm.test.tester.test_pipeline import EthosU55PipelineINT
+
+input_t = Tuple[torch.Tensor]
+
+
+class PerChannelAffineRank4(torch.nn.Module):
+    """Per-channel scale-mul + bias-add with a transpose to engage
+    propagation.
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.scale = torch.nn.Parameter(torch.rand(240) * 0.1 + 0.9)
+        self.bias = torch.nn.Parameter(torch.randn(240) * 0.1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        y = x.reshape(1, 1, 1, x.shape[-1])
+        y = y * self.scale + self.bias
+        y = y.permute(0, 3, 1, 2)
+        y = y.reshape(1, x.shape[-1])
+        y = y * self.scale + self.bias
+        return y
+
+    data = (torch.randn(1, 240),)
+
+
+def test_per_channel_affine_rank4_u55_INT() -> None:
+    pipeline = EthosU55PipelineINT[input_t](
+        PerChannelAffineRank4(),
+        PerChannelAffineRank4.data,
+        aten_ops=[],
+        exir_ops=[],
+        run_on_fvp=False,
+        a16w8_quantization=True,
+    )
+    pipeline.run()