Add var decomposition for Arm backend

- Add unittests for var, var.dim and var.correction
- Add decmposition for var.correction. var and var.dim are both
converted to var.correction earlier in the lowering.

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

Author: oscarandersson8218

backends/arm/_passes/arm_pass_manager.py

@@ -20,6 +20,7 @@
 )
 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.decompose_var_pass import DecomposeVarPass
 from executorch.backends.arm._passes.insert_squeeze_after_sum_pass import (
     InsertSqueezeAfterSumPass,
 )
@@ -49,6 +50,7 @@ def transform_to_backend_pipeline(
         self.add_pass(SizeAdjustConv2DPass())
         self.add_pass(RemoveClonePass())
         self.add_pass(ConvertExpandCopyToRepeatPass())
+        self.add_pass(DecomposeVarPass())
         self.add_pass(ConvertMeanDimToAveragePool())
         self.add_pass(DecomposeMeanDimPass())
         self.add_pass(DecomposeDivPass())
@@ -63,6 +65,8 @@ def transform_to_backend_pipeline(
         return self._transform(exported_program.graph_module)
 
     def transform_for_annotation_pipeline(self, graph_module: torch.fx.GraphModule):
+        self.add_pass(DecomposeVarPass())
+        self.add_pass(DecomposeMeanDimPass())
         self.add_pass(DecomposeDivPass())
         self.add_pass(ScalarsToAttributePass())
         self.add_pass(DecomposeMeanDimPass())

backends/arm/_passes/decompose_var_pass.py

@@ -0,0 +1,83 @@
+# 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_var_decomposition(op) -> tuple:
+    if op == exir_ops.edge.aten.var.correction:
+        return (
+            exir_ops.edge.aten.mean.dim,
+            exir_ops.edge.aten.sub.Tensor,
+            exir_ops.edge.aten.mul.Tensor,
+            exir_ops.edge.aten.sum.dim_IntList,
+            exir_ops.edge.aten.full.default,
+        )
+    if op in (torch.ops.aten.var.correction, torch.ops.aten.var.dim):
+        return (
+            torch.ops.aten.mean.dim,
+            torch.ops.aten.sub.Tensor,
+            torch.ops.aten.mul.Tensor,
+            torch.ops.aten.sum.dim_IntList,
+            torch.ops.aten.full,
+        )
+    raise RuntimeError(f"Can't get var decomposition for op {op}")
+
+
+class DecomposeVarPass(ExportPass):
+    """
+    This pass decomposes var.correction and var.dim into smaller ops (see https://pytorch.org/docs/stable/generated/torch.var.html)
+
+    Example:
+        y = var_correction(x, dim, keepdim, correction)
+    Becomes:
+        mean = mean(x, dim)
+        diff = sub(x, mean)
+        squared_diff = mul(diff, diff)
+        sum = sum(squared_diff, dim)
+        y = div(sum, max(0, N-correction))
+    """
+
+    def call_operator(self, op, args, kwargs, meta):
+        if op not in (
+            exir_ops.edge.aten.var.correction,
+            torch.ops.aten.var.correction,
+            torch.ops.aten.var.dim,
+        ):
+            return super().call_operator(op, args, kwargs, meta)
+        shape = meta["val"].size()
+        dtype = meta["val"].dtype
+        dim = args[1] if len(args) > 1 else list(range(len(shape)))
+        if op == torch.ops.aten.var.dim:
+            correction = args[-2]
+            keepdim = args[-1]
+        else:
+            correction = kwargs["correction"]
+            keepdim = kwargs.get("keepdim", False)
+        if not keepdim:
+            return super().call_operator(op, args, kwargs, meta)
+
+        x = args[0]
+        input_shape = x.data.size()
+        N = 1
+        for d in dim:
+            N *= input_shape[d]
+
+        mean_op, diff_op, mul_op, sum_op, full_op = get_var_decomposition(op)
+        mean = super().call_operator(mean_op, (x, dim, keepdim), {}, meta)
+        diff = super().call_operator(diff_op, (x, mean), {}, meta)
+        squared_diff = super().call_operator(mul_op, (diff, diff), {}, meta)
+        sum = super().call_operator(sum_op, (squared_diff, dim, keepdim), {}, meta)
+        full = super().call_operator(
+            full_op,
+            ([1 for _ in shape], 1 / max(0, N - correction)),
+            {"dtype": dtype},
+            meta,
+        )
+        return super().call_operator(mul_op, (sum, full), {}, meta)

backends/arm/arm_partitioner.py

@@ -67,6 +67,7 @@ def is_node_supported(self, submodules, node: torch.fx.Node) -> bool:
             exir_ops.edge.aten.view_copy.default,
             exir_ops.edge.aten.clone.default,
             exir_ops.edge.aten.mean.dim,
+            exir_ops.edge.aten.var.correction,
             exir_ops.edge.aten.unsqueeze_copy.default,
             exir_ops.edge.aten.squeeze_copy.dims,
             operator.getitem,
@@ -87,6 +88,9 @@ def is_node_supported_custom(self, node: torch.fx.Node) -> bool:
         if node.target == exir_ops.edge.aten.mean.dim:
             keep_dim = node.args[2] if len(node.args) > 2 else False
             return keep_dim
+        if node.target == exir_ops.edge.aten.var.correction:
+            keep_dim = node.kwargs.get("keepdim", False)
+            return keep_dim
         return True
 
 

backends/arm/operators/op_full.py

@@ -52,5 +52,7 @@ def define_node(
             dtype = ts.DType.FP32
             data = np.full(shape, value, dtype=np.float32)
 
-        tosa_graph.addConst(shape, dtype, data, "full-const")
-        tosa_graph.addOperator(ts.TosaOp.Op.IDENTITY, ["full-const"], [output.name])
+        tosa_graph.addConst(shape, dtype, data, node.name + "full-const")
+        tosa_graph.addOperator(
+            ts.TosaOp.Op.IDENTITY, [node.name + "full-const"], [output.name]
+        )

backends/arm/quantizer/quantization_annotation/sub_annotator.py

@@ -6,8 +6,6 @@
 
 # pyre-unsafe
 
-import itertools
-import operator
 from typing import Callable, List, Optional
 
 import torch
@@ -16,7 +14,6 @@
 from executorch.backends.arm.quantizer.quantization_config import QuantizationConfig
 from torch.ao.quantization.quantizer import QuantizationAnnotation
 from torch.fx import GraphModule, Node
-from torch.fx.passes.utils.source_matcher_utils import get_source_partitions
 
 
 @register_annotator("sub")
@@ -25,14 +22,12 @@ def _annotate_sub(
     quantization_config: QuantizationConfig,
     filter_fn: Optional[Callable[[Node], bool]] = None,
 ) -> Optional[List[List[Node]]]:
-    sub_partitions = get_source_partitions(
-        gm.graph, [operator.sub, torch.sub, operator.isub], filter_fn
-    )
-    sub_partitions = list(itertools.chain.from_iterable(sub_partitions.values()))
     annotated_partitions = []
-    for sub_partition in sub_partitions:
-        annotated_partitions.append(sub_partition.nodes)
-        sub_node = sub_partition.output_nodes[0]
+    for node in gm.graph.nodes:
+        if node.target not in (torch.ops.aten.sub.Tensor,):
+            continue
+        annotated_partitions.append(node)
+        sub_node = node
         if arm_quantizer_utils.is_annotated(sub_node):
             continue