Add meandim decomposition

Signed-off-by: Oscar Andersson <[email protected]>
Change-Id: I6e9b3da1384d33903b039696d6c156b1a2f1841a

Author: oscarandersson8218

backends/arm/_passes/arm_pass_manager.py

@@ -19,6 +19,7 @@
     ConvertSplitToSlicePass,
 )
 from executorch.backends.arm._passes.decompose_div_pass import DecomposeDivPass
+from executorch.backends.arm._passes.decompose_meandim_pass import DecomposeMeanDimPass
 from executorch.backends.arm._passes.insert_squeeze_after_sum_pass import (
     InsertSqueezeAfterSumPass,
 )
@@ -49,6 +50,7 @@ def transform_to_backend_pipeline(
         self.add_pass(RemoveClonePass())
         self.add_pass(ConvertExpandCopyToRepeatPass())
         self.add_pass(ConvertMeanDimToAveragePool())
+        self.add_pass(DecomposeMeanDimPass())
         self.add_pass(DecomposeDivPass())
         self.add_pass(InsertSqueezeAfterSumPass())
         self.add_pass(ConvertSplitToSlicePass())
@@ -63,4 +65,5 @@ def transform_to_backend_pipeline(
     def transform_for_annotation_pipeline(self, graph_module: torch.fx.GraphModule):
         self.add_pass(DecomposeDivPass())
         self.add_pass(ScalarsToAttributePass())
+        self.add_pass(DecomposeMeanDimPass())
         return self._transform(graph_module)

backends/arm/_passes/decompose_meandim_pass.py

@@ -0,0 +1,66 @@
+# Copyright 2024 Arm Limited and/or its 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.
+
+import torch
+from executorch.exir.dialects._ops import ops as exir_ops
+from executorch.exir.pass_base import ExportPass
+
+
+def get_meandim_decomposition(op) -> tuple:
+    if op == exir_ops.edge.aten.mean.dim:
+        return (
+            exir_ops.edge.aten.sum.dim_IntList,
+            exir_ops.edge.aten.full.default,
+            exir_ops.edge.aten.mul.Tensor,
+        )
+    if op == torch.ops.aten.mean.dim:
+        return (
+            torch.ops.aten.sum.dim_IntList,
+            torch.ops.aten.full.default,
+            torch.ops.aten.mul.Tensor,
+        )
+    raise RuntimeError(f"Can't get meandim decomposition for op {op}")
+
+
+class DecomposeMeanDimPass(ExportPass):
+    """
+    This pass decomposes meandim into a sum and mul node.
+
+    Example:
+        y = mean_dim(x, dim, keepdim)
+    Becomes:
+        sum = sum.dim_IntList(x, dim, keepdim)
+        y = mul(sum, 1/N)
+    """
+
+    def call_operator(self, op, args, kwargs, meta):
+        if op not in (exir_ops.edge.aten.mean.dim, torch.ops.aten.mean.dim):
+            return super().call_operator(op, args, kwargs, meta)
+
+        x = args[0]
+        dim = args[1]
+        keepdim = args[2] if len(args) > 2 else False
+        if not keepdim:
+            return super().call_operator(op, args, kwargs, meta)
+        # if keepdim == True and dim == [-1, -2], mean.dim can be
+        # decomposed to avg_pool2d. This is handled by ConvertMeanDimToAveragePool.
+        if dim == [-1, -2]:
+            # Simply return the mean.dim operator for future decomposition.
+            return super().call_operator(op, args, kwargs, meta)
+        shape = meta["val"].size()
+        dtype = meta["val"].dtype
+        input_shape = x.data.size()
+        N = 1
+        for d in dim:
+            N *= input_shape[d]
+
+        sum_op, full_op, mul_op = get_meandim_decomposition(op)
+
+        sum = super().call_operator(sum_op, (x, dim, keepdim), {}, meta)
+        full = super().call_operator(
+            full_op, ([1] * len(shape), 1 / N), {"dtype": dtype}, meta
+        )
+        return super().call_operator(mul_op, (sum, full), {}, meta)

backends/arm/arm_partitioner.py

@@ -85,10 +85,8 @@ def is_node_supported(self, submodules, node: torch.fx.Node) -> bool:
 
     def is_node_supported_custom(self, node: torch.fx.Node) -> bool:
         if node.target == exir_ops.edge.aten.mean.dim:
-            dim = node.args[1]
-            keep_dim = node.args[2]
-            if dim != [-1, -2] or keep_dim is False:
-                return False
+            keep_dim = node.args[2] if len(node.args) > 2 else False
+            return keep_dim
         return True
 
 

backends/arm/operators/__init__.py

@@ -20,7 +20,6 @@
     op_get_item,
     op_hardtanh,
     op_log,
-    op_mean_dim,
     op_mm,
     op_mul,
     op_permute,