Arm backend: Replace squeeze with view in DecomposeAnyPass (#15733)

Rename ConvertAnyDefaultDimDimsPass to DecomposeAnyPass because, in
practice, it does same thing for the any op as DecomposeSumPass does for
sum.

When keepdim=False, create a single view_copy to the reduced shape
instead of inserting squeeze_copy. This eliminates the chronological
dependency of ConvertSqueezesToViewPass.

cc @freddan80 @per @zingo @oscarandersson8218 @digantdesai

Signed-off-by: Martin Lindström <[email protected]>
Co-authored-by: Martin Lindström <[email protected]>

Author: martinlsm

backends/arm/_passes/__init__.py

@@ -13,7 +13,6 @@
 from .cast_int64_pass import CastInt64BuffersToInt32Pass  # noqa
 from .cast_to_int32_pass import CastToInt32Pass  # noqa
 from .conv1d_unsqueeze_pass import Conv1dUnsqueezePass  # noqa
-from .convert_any_default_dim_dims_pass import ConvertAnyDefaultDimDimsPass  # noqa
 from .convert_elu_params import ConvertELUParamsPass  # noqa
 from .convert_expand_copy_to_repeat import ConvertExpandCopyToRepeatPass  # noqa
 from .convert_full_like_to_full_pass import ConvertFullLikeToFullPass  # noqa
@@ -31,6 +30,7 @@
 from .decompose_adaptive_avg_pool2d_pass import DecomposeAdaptiveAvgPool2dPass  # noqa
 from .decompose_add_sub_alpha_pass import DecomposeAddSubAlphaPass  # noqa
 from .decompose_addmm_pass import DecomposeAddmmPass  # noqa
+from .decompose_any_pass import DecomposeAnyPass  # noqa
 from .decompose_asin_and_acos_pass import DecomposeAsinAndAcosPass  # noqa
 from .decompose_asinh_pass import DecomposeAsinhPass  # noqa
 from .decompose_atan_pass import DecomposeAtanPass  # noqa

backends/arm/_passes/arm_pass_manager.py

@@ -18,7 +18,6 @@
     CastToInt32Pass,
     ComputeConstantOpsAOT,
     Conv1dUnsqueezePass,
-    ConvertAnyDefaultDimDimsPass,
     ConvertELUParamsPass,
     ConvertExpandCopyToRepeatPass,
     ConvertFullLikeToFullPass,
@@ -35,6 +34,7 @@
     DecomposeAdaptiveAvgPool2dPass,
     DecomposeAddmmPass,
     DecomposeAddSubAlphaPass,
+    DecomposeAnyPass,
     DecomposeAsinAndAcosPass,
     DecomposeAsinhPass,
     DecomposeAtanhPass,
@@ -241,7 +241,7 @@ def _tosa_pipeline(
         self.add_pass(DecomposeDivPass())
         self.add_pass(DecomposeSoftmaxPass())
         self.add_pass(ConvertMinMaxPass())
-        self.add_pass(ConvertAnyDefaultDimDimsPass())
+        self.add_pass(DecomposeAnyPass())
         self.add_pass(DecomposeAdaptiveAvgPool2dPass())
         self.add_pass(DecomposeAvgPool2d())
         self.add_pass(

…ses/convert_any_default_dim_dims_pass.py …ckends/arm/_passes/decompose_any_pass.pybackends/arm/_passes/convert_any_default_dim_dims_pass.py renamed to backends/arm/_passes/decompose_any_pass.py backends/arm/_passes/convert_any_default_dim_dims_pass.py renamed to backends/arm/_passes/decompose_any_pass.py

@@ -7,9 +7,7 @@
 
 import torch
 from executorch.backends.arm._passes import ArmPass
-from executorch.backends.arm._passes.convert_squeezes_to_view import (
-    ConvertSqueezesToViewPass,
-)
+from executorch.backends.arm._passes.arm_pass_utils import get_first_fake_tensor
 from executorch.exir.dialects._ops import (  # type: ignore[import-not-found]
     ops as exir_ops,
 )
@@ -19,38 +17,40 @@
 )
 
 
-class ConvertAnyDefaultDimDimsPass(ArmPass):
+class DecomposeAnyPass(ArmPass):
     """
-    Converts any.default, any.dim and any.dims to a sequence of any.dim by unrolling multi-dimensional reduction.
-    Please refer to KeepDimsFalseToSqueezePass for an explanation of this coversion.
+    Converts any.default, any.dim and any.dims to a sequence of any.dim by
+    unrolling multi-dimensional reductions with keepdim=True. If keepdim=False
+    was requested, the final shape adjustment is implemented with a
+    view_copy.default to the reduced shape.
 
     Example 1
     Original:
-        any()  # x.shape: [dim1, dim2, ..., dimn]
+        any.dim()  # x.shape: [dim1, dim2, ..., dimn]
     After pass:
         any.dim(dim1, keepdim = True)
         any.dim(dim2, keepdim = True)
         ...
         any.dim(dimn, keepdim = True)
-        squeeze(dim = [dim1, dim2, ...., dimn])
+        view_copy(shape = squeezed_shape)
 
     Example 2
     Original:
         any.dim(dim1, keepdim = False)
     After pass:
         any.dim(dim1, keepdim = True)
-        squeeze(dim = [dim1])
+        view_copy(shape = squeezed_shape)
 
     Example 3
     Original:
         any.dims([dim1, dim2], keepdim = False)
     After pass:
         any.dim(dim1, keepdim = True)
         any.dim(dim2, keepdim = True)
-        squeeze(dim = [dim1, dim2])
+        view_copy(shape = squeezed_shape)
     """
 
-    _passes_required_after: Set[Type[ExportPass]] = {ConvertSqueezesToViewPass}
+    _passes_required_after: Set[Type[ExportPass]] = set()
 
     def call(self, graph_module: torch.fx.GraphModule):
         modified = False
@@ -67,40 +67,40 @@ def call(self, graph_module: torch.fx.GraphModule):
             if len(node.args) == 1:
                 # any.default(input)
                 input_node = (node.args)[0]
-                dims = range(len(input_node.meta["val"].shape))
+                dims_to_reduce = range(len(input_node.meta["val"].shape))
                 keepdim = False
             elif len(node.args) == 2:
                 # any.dim/dims(input, dims=dims)
-                input_node, dims = node.args
+                input_node, dims_to_reduce = node.args
                 keepdim = False
             elif len(node.args) == 3:
                 # any.dim/dims(input, dims=dims, keepdim=keepdim)
-                input_node, dims, keepdim = node.args
+                input_node, dims_to_reduce, keepdim = node.args
             else:
                 raise RuntimeError(
                     f"Unexpected arg size {len(node.args)} in {node.name}"
                 )
             try:
-                iter(dims)
+                iter(dims_to_reduce)
             except:
-                dims = [dims]  # type: ignore[assignment]
+                dims_to_reduce = [dims_to_reduce]  # type: ignore[assignment]
             else:
-                dims = list(dims)  # type: ignore[assignment]
+                dims_to_reduce = list(dims_to_reduce)  # type: ignore[assignment]
 
             # Unroll multi-dimensional reduction and keep-dims arg
             with graph_module.graph.inserting_before(node):
-                for dim in dims:
+                for dim in dims_to_reduce:
                     args = (input_node, dim, True)
                     input_node = graph_module.graph.create_node(
                         "call_function", exir_ops.edge.aten.any.dim, args, node.kwargs
                     )
 
                 if not keepdim:
-                    args = (input_node, dims)  # type: ignore[assignment]
+                    output_shape = list(get_first_fake_tensor(node).shape)
                     input_node = graph_module.graph.create_node(
                         "call_function",
-                        exir_ops.edge.aten.squeeze_copy.dims,
-                        args,
+                        exir_ops.edge.aten.view_copy.default,
+                        (input_node, output_shape),
                     )
 
             node.replace_all_uses_with(input_node)

backends/arm/operators/op_any.py

@@ -46,7 +46,7 @@ def define_node(
         )  # process the negative index
         keep_dim = cast(bool, inputs[2].number if len(inputs) > 2 else False)
         if not keep_dim:
-            raise ValueError("This case should be handled by ConvertAnyDimDimsPass")
+            raise ValueError("This case should be handled by DecomposeAnyPass")
 
         attr = ts.TosaSerializerAttribute()
         attr.ReduceAnyAttribute(inputs[0].dim_order.index(dim))