Arm backend: Add layer_norm decomposition (#6288)

Summary:
Add decomposition for layer_norm, mean_dim and var_dim.

Pull Request resolved: #6288

Reviewed By: cccclai

Differential Revision: D64473056

Pulled By: digantdesai

fbshipit-source-id: a3d4b0fbb0ea1d56b372b617f96dcbab41f72bb7

Author: oscarandersson8218

backends/arm/_passes/arm_pass_manager.py

@@ -19,6 +19,11 @@
     ConvertSplitToSlicePass,
 )
 from executorch.backends.arm._passes.decompose_div_pass import DecomposeDivPass
+from executorch.backends.arm._passes.decompose_layernorm_pass import (
+    DecomposeLayerNormPass,
+)
+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,
 )
@@ -53,7 +58,10 @@ def transform_to_backend_pipeline(
         self.add_pass(SizeAdjustConv2DPass())
         self.add_pass(RemoveClonePass())
         self.add_pass(ConvertExpandCopyToRepeatPass())
+        self.add_pass(DecomposeLayerNormPass())
+        self.add_pass(DecomposeVarPass())
         self.add_pass(ConvertMeanDimToAveragePool())
+        self.add_pass(DecomposeMeanDimPass())
         self.add_pass(MatchArgRanksPass(exported_program))
         self.add_pass(DecomposeDivPass())
         self.add_pass(InsertSqueezeAfterSumPass())
@@ -67,6 +75,9 @@ 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(DecomposeLayerNormPass())
+        self.add_pass(DecomposeVarPass())
+        self.add_pass(DecomposeMeanDimPass())
         self.add_pass(ScalarsToAttributePass())
         self.add_pass(DecomposeDivPass())
         return self._transform(graph_module)

backends/arm/_passes/decompose_layernorm_pass.py

@@ -0,0 +1,152 @@
+# 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 operator
+
+import torch
+from executorch.backends.arm._passes.arm_pass_utils import create_node
+from executorch.exir.dialects._ops import ops as exir_ops
+from executorch.exir.pass_base import ExportPass, PassResult
+
+
+def get_layer_norm_decomposition(op) -> tuple:
+    if op == exir_ops.edge.aten.native_layer_norm.default:
+        return (
+            exir_ops.edge.aten.mean.dim,
+            exir_ops.edge.aten.sub.Tensor,
+            exir_ops.edge.aten.var.correction,
+            exir_ops.edge.aten.full.default,
+            exir_ops.edge.aten.add.Tensor,
+            exir_ops.edge.aten.rsqrt.default,
+            exir_ops.edge.aten.mul.Tensor,
+            exir_ops.edge.aten.view_copy.default,
+        )
+    if op == torch.ops.aten.layer_norm.default:
+        return (
+            torch.ops.aten.mean.dim,
+            torch.ops.aten.sub.Tensor,
+            torch.ops.aten.var.correction,
+            torch.ops.aten.full.default,
+            torch.ops.aten.add.Tensor,
+            torch.ops.aten.rsqrt.default,
+            torch.ops.aten.mul.Tensor,
+            torch.ops.aten.view_copy.default,
+        )
+    raise RuntimeError(f"Can't get layer_norm composition for op {op}")
+
+
+class DecomposeLayerNormPass(ExportPass):
+    """
+    layernorm is defined as: ((x - E[x]) / sqrt(Var[x] + eps)) * weights + bias
+    Decompose layernorm(x, normalized_shape, weights, bias, eps) to a sequence of:
+    mean        = op_mean(x, dims)           # E[x]
+    var         = op_var(x, dims)            # Var[x]
+    denominator = op_sub(x, mean)            # (x - E[x])
+    add         = op_add(var, eps)           # Var[x] + eps
+    rsqrt       = op_rsqrt(add)              # 1 / sqrt(Var[x] + eps)
+    mul         = op_mul(denominator, rsqrt) # ((x - E[x]) / sqrt(Var[x] + eps)) * weigths
+    bias        = op_add(mul, bias)          # ((x - E[x]) / sqrt(Var[x] + eps)) * weigths + bias
+
+    Source: https://pytorch.org/docs/stable/generated/torch.nn.LayerNorm.html
+    """
+
+    def call(self, graph_module: torch.fx.GraphModule):
+        for node in graph_module.graph.nodes:
+            if node.op != "call_function" or node.target not in (
+                exir_ops.edge.aten.native_layer_norm.default,
+                torch.ops.aten.layer_norm.default,
+            ):
+                continue
+
+            # epsilon default value
+            epsilon = torch.finfo().eps
+            weights = None
+            bias = None
+            args = node.args
+            meta = node.meta
+            match len(args):
+                case 5:
+                    x, normalized_shape, weights, bias, epsilon = args
+                case 4:
+                    x, normalized_shape, weights, bias = args
+                case 3:
+                    x, normalized_shape, weights = args
+                case _:
+                    x, normalized_shape = args
+
+            n_dims = len(normalized_shape)
+            if isinstance(meta["val"], tuple):
+                shape = meta["val"][0].size()
+            else:
+                shape = meta["val"].size()
+            dtype = meta["val"][0].dtype
+            rank = len(shape)
+            dims = list(range(-1, -1 * (n_dims + 1), -1))
+            dims = [dim % rank for dim in dims]
+            weights_reshaped_shape = [shape[i] if i in dims else 1 for i in range(rank)]
+            epsilon_reshaped_shape = [1] * rank
+
+            (
+                mean_op,
+                sub_op,
+                var_op,
+                full_op,
+                add_op,
+                rsqrt_op,
+                mul_op,
+                view_op,
+            ) = get_layer_norm_decomposition(node.target)
+            with graph_module.graph.inserting_before(node):
+                keepdim = True
+                mean = create_node(graph_module.graph, mean_op, args=(x, dims, keepdim))
+                sub = create_node(graph_module.graph, sub_op, args=(x, mean))
+                var = create_node(
+                    graph_module.graph,
+                    var_op,
+                    args=(x, dims),
+                    kwargs={"correction": 0, "keepdim": keepdim},
+                )
+                full = create_node(
+                    graph_module.graph,
+                    full_op,
+                    args=(epsilon_reshaped_shape, epsilon),
+                    kwargs={"dtype": dtype},
+                )
+                add0 = create_node(graph_module.graph, add_op, args=(var, full))
+                rsqrt = create_node(graph_module.graph, rsqrt_op, args=(add0,))
+                mul0 = create_node(graph_module.graph, mul_op, args=(sub, rsqrt))
+                if weights is not None:
+                    weights_reshaped = create_node(
+                        graph_module.graph,
+                        view_op,
+                        args=(weights, weights_reshaped_shape),
+                    )
+                    mul1 = create_node(
+                        graph_module.graph, mul_op, args=(mul0, weights_reshaped)
+                    )
+                else:
+                    mul1 = mul0
+                output = mul1
+                if bias is not None:
+                    bias_reshaped_shape = weights_reshaped_shape
+                    bias_reshaped = create_node(
+                        graph_module.graph, view_op, args=(bias, bias_reshaped_shape)
+                    )
+                    output = create_node(
+                        graph_module.graph, add_op, args=(mul1, bias_reshaped)
+                    )
+
+                users = [user for user in node.users if node != user]
+                node.replace_all_uses_with(output)
+                for user in users:
+                    if user.target == operator.getitem:
+                        user.replace_all_uses_with(output)
+                graph_module.graph.erase_node(node)
+                graph_module.graph.eliminate_dead_code()
+        graph_module.recompile()
+        graph_module = super().call(graph_module).graph_module
+
+        return PassResult(graph_module, True)

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/_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

@@ -8,7 +8,7 @@
 import logging
 import operator
 import os
-from typing import final, List
+from typing import cast, final, List
 
 import torch
 from executorch.backends.arm.arm_backend import ArmBackend  # usort: skip
@@ -53,6 +53,7 @@ def is_node_supported(self, submodules, node: torch.fx.Node) -> bool:
             exir_ops.edge.aten.full.default,
             exir_ops.edge.aten.mul.Tensor,
             exir_ops.edge.aten._native_batch_norm_legit_no_training.default,
+            exir_ops.edge.aten.native_layer_norm.default,
             exir_ops.edge.aten.avg_pool2d.default,
             exir_ops.edge.aten.sigmoid.default,
             exir_ops.edge.aten.mm.default,
@@ -67,6 +68,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,
@@ -85,10 +87,11 @@ 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 cast(bool, keep_dim)
+        if node.target == exir_ops.edge.aten.var.correction:
+            keep_dim = node.kwargs.get("keepdim", False)
+            return cast(bool, 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,

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]
+        )