Handle rank-changing views in RemovePermutesAroundElementwiseOps (pytorch#19538)

Differential Revision: D104775244

Pull Request resolved: pytorch#19538

Author: mcremon-meta

backends/arm/_passes/remove_permutes_around_elementwise_tosa_ops.py

@@ -11,13 +11,15 @@
 
 
 class RemovePermutesAroundElementwiseTosaOps(RemovePermutesAroundElementwiseOps):
-    permutable_ops = {
-        *RemovePermutesAroundElementwiseOps.permutable_ops,
-        *TableOps.unary_table_ops.keys(),
-        *TableOps.special_table_ops,
-        exir_ops.backend.tosa.RESCALE.default,
-        exir_ops.backend.tosa.TABLE.default,
-    }
+    def __init__(self) -> None:
+        super().__init__(
+            extra_permutable_ops={
+                *TableOps.unary_table_ops.keys(),
+                *TableOps.special_table_ops,
+                exir_ops.backend.tosa.RESCALE.default,
+                exir_ops.backend.tosa.TABLE.default,
+            }
+        )
 
     def permute_subgraph(self, subgraph):
         # Original function will always permute constant nodes which is wrong for table ops

backends/cadence/aot/remove_ops.py

@@ -603,16 +603,16 @@ def maybe_remove_or_replace(self, node: Node) -> bool:
 
 @register_cadence_pass(CadencePassAttribute(opt_level=2))
 class RemovePermutesAroundElementwiseOps(_SharedRemovePermutesAroundElementwiseOps):
-    permutable_ops: set[EdgeOpOverload] = (
-        _SharedRemovePermutesAroundElementwiseOps.permutable_ops
-        | {
-            exir_ops.edge.cadence.quantize_per_tensor.default,
-            exir_ops.edge.cadence.dequantize_per_tensor.default,
-            exir_ops.edge.cadence.quantized_relu.per_tensor,
-            exir_ops.edge.cadence.requantize.per_tensor,
-            exir_ops.edge.cadence.quantized_add.per_tensor,
-        }
-    )
+    def __init__(self) -> None:
+        super().__init__(
+            extra_permutable_ops={
+                exir_ops.edge.cadence.quantize_per_tensor.default,
+                exir_ops.edge.cadence.dequantize_per_tensor.default,
+                exir_ops.edge.cadence.quantized_relu.per_tensor,
+                exir_ops.edge.cadence.requantize.per_tensor,
+                exir_ops.edge.cadence.quantized_add.per_tensor,
+            }
+        )
 
 
 @register_cadence_pass(CadencePassAttribute(opt_level=2))