Merge branch 'main' into Meta_fusing_conv1d_relu_residual_add

Author: digantdesai

backends/arm/README.md

@@ -209,6 +209,38 @@ List of model specific and optional passes:
 - InsertCastForOpsWithInt64InputPass
     - Functionality:
         - For LLMs such as LLama, some opeartors like aten.embedding have int64 input. In order to lower these operators to TOSA, this pass will insert a casting node that converts the input from int64 to int32.
-        - Example usage: backends/arm/test/models/test_llama.py
     - Supported Ops:
         - aten.embedding.default, aten.slice_copy.Tensor
+    - Example usage:
+        - backends/arm/test/models/test_llama.py
+
+- ConvertInt64ConstOpsToInt32Pass
+    - Functionalities:
+      - Rewrites constant-producing ops that output int64 to instead output int32, when values are within int32 bounds.
+    - Supported Ops:
+      - `torch.full`, `torch.arange`, `torch.eye`, `torch.linspace`, `torch.tensor`
+    - Example usage:
+        - backends/arm/test/models/stable_diffusion/test_CLIPTextModelWithProjection.py
+        - backends/arm/test/models/stable_diffusion/test_T5EncoderModel.py
+
+- ConvertInt64OutputOpsToInt32Pass
+    - Overview:
+      - Rewrites or removes operations that produce int64 outputs, converting them to int32 where possible.
+      - Overflow checks are applied selectively; for ops without such checks, users need to ensure values fit within the int32 range.
+    - Functionalities:
+        1. Handling casting to int64:
+            - (1) int32 -> int64:
+                - Removes the cast and redirect uses of int64 to int32
+            - (2) other types -> int64:
+                - Rewrites the cast to other types -> int32
+            - Supported Ops:
+              - torch.ops.aten.to.\[dtype|dtype_layout\]
+              - exir_ops.edge.dim_order_ops._to_dim_order_copy.default
+        2. Post-process argmax outputs:
+            - Inserts an int64->int32 cast after the argmax operations that produce int64 outputs:
+            - Supported Ops:
+              - torch.ops.aten.argmax.default
+              - exir_ops.edge.aten.argmax.default
+    - Example usage:
+      - (Functionality 1) backends/arm/test/models/stable_diffusion/test_T5EncoderModel.py
+      - (Functionality 2) backends/arm/test/models/stable_diffusion/test_CLIPTextModelWithProjection.py

backends/arm/_passes/__init__.py

@@ -14,8 +14,11 @@
 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
+from .convert_int64_const_ops_to_int32 import ConvertInt64ConstOpsToInt32Pass  # noqa
+from .convert_int64_output_ops_to_int32 import ConvertInt64OutputOpsToInt32Pass  # noqa
 from .convert_int_pow_to_mul import ConvertIntPowToMuls  # noqa
 from .convert_minmax_pass import ConvertMinMaxPass  # noqa
 from .convert_split_to_slice import ConvertSplitToSlicePass  # noqa
@@ -34,6 +37,7 @@
 from .decompose_cosine_similarity_pass import DecomposeCosineSimilarityPass  # noqa
 from .decompose_cumsum_pass import DecomposeCumsumPass  # noqa
 from .decompose_div_pass import DecomposeDivPass  # noqa
+from .decompose_elu_pass import DecomposeEluPass  # noqa
 from .decompose_embedding_pass import DecomposeEmbeddingPass  # noqa  # noqa
 from .decompose_expm1_pass import DecomposeExpm1Pass  # noqa
 from .decompose_gelu_pass import DecomposeGeluPass  # noqa

backends/arm/_passes/arm_pass_manager.py

@@ -18,8 +18,11 @@
     ComputeConstantOpsAOT,
     Conv1dUnsqueezePass,
     ConvertAnyDefaultDimDimsPass,
+    ConvertELUParamsPass,
     ConvertExpandCopyToRepeatPass,
     ConvertFullLikeToFullPass,
+    ConvertInt64ConstOpsToInt32Pass,
+    ConvertInt64OutputOpsToInt32Pass,
     ConvertIntPowToMuls,
     ConvertMinMaxPass,
     ConvertMmToBmmPass,
@@ -39,6 +42,7 @@
     DecomposeCosineSimilarityPass,
     DecomposeCumsumPass,
     DecomposeDivPass,
+    DecomposeEluPass,
     DecomposeEmbeddingPass,
     DecomposeExpm1Pass,
     DecomposeGeluPass,
@@ -98,6 +102,7 @@
 from executorch.backends.transforms.remove_getitem_op import RemoveGetItemPass
 from executorch.exir import ExportedProgram
 from executorch.exir.pass_manager import PassManager
+from executorch.exir.passes.remove_graph_asserts_pass import RemoveGraphAssertsPass
 from torch.fx import GraphModule
 
 
@@ -132,6 +137,7 @@ def _tosa_INT_pipeline(self, exported_program: ExportedProgram) -> GraphModule:
         self.add_pass(ReplaceScalarWithTensorArgPassTOSABI())
         self.add_pass(AnnotateDecomposedMatmulPass())
         self.add_pass(QuantizeOperatorArguments())
+        self.add_pass(ConvertELUParamsPass())
         self.add_pass(FoldAndAnnotateQParamsPass(exported_program))  # type: ignore[call-arg]
         self.add_pass(RetraceFoldedDtypesPass())
         self.add_pass(UnsqueezeScalarPlaceholdersPass(exported_program))
@@ -180,6 +186,8 @@ def _tosa_FP_pipeline(self, exported_program: ExportedProgram) -> GraphModule:
         self.add_pass(DecomposeAtanPass())
         self.add_pass(DecomposeAtanhPass())
         self.add_pass(DecomposeAddmmPass())
+        self.add_pass(DecomposeEluPass())
+        self.add_pass(DecomposeExpm1Pass())
         self.add_pass(ConvertIntPowToMuls())
         self.add_pass(CastBoolToInt8Pass())
         self.add_pass(DecomposeSinhPass())
@@ -258,6 +266,11 @@ def transform_to_backend_pipeline(self, exported_program: ExportedProgram):
             )
 
     def transform_for_annotation_pipeline(self, graph_module: GraphModule):
+        self.add_pass(
+            RemoveGraphAssertsPass()
+        )  # ConvertInt64ConstOpsToInt32Pass requires this pass to remove the assertation in Graph
+        self.add_pass(ConvertInt64ConstOpsToInt32Pass())
+        self.add_pass(ConvertInt64OutputOpsToInt32Pass())
         self.add_pass(InsertCastForOpsWithInt64InputPass())
         self.add_pass(DecomposeEmbeddingPass())
         self.add_pass(DecomposeScaledDotProductAttention())

backends/arm/_passes/convert_elu_params.py

@@ -0,0 +1,53 @@
+# Copyright 2025 Arm Limited and/or its affiliates.
+#
+# 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.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
+
+
+class ConvertELUParamsPass(ExportPass):
+    """
+    Pass to convert the input_scale kwarg of ELU operator from float to
+    int.
+
+    It has been set to 2 as the outputs seem to stay the same regardless of what
+    the value of input_scale is, as long as that value is not 1.
+    """
+
+    def call(self, graph_module: torch.fx.GraphModule):
+        modified_graph = False
+        graph = graph_module.graph
+        node_list = graph.find_nodes(
+            op="call_function", target=exir_ops.edge.aten.elu.default
+        )
+        for node in node_list:
+            with graph.inserting_after(node):
+                replace_node = create_node(graph, exir_ops.edge.aten.elu.default)
+                old_args = list(node.args)
+
+                alpha = old_args[1] if len(old_args) > 1 else 1.0
+                scale = 1.0
+                input_scale = 2.0
+
+                replace_node.args = (old_args[0],)
+
+                updated_kwargs = dict(node.kwargs)
+                updated_kwargs["alpha"] = int(alpha)
+                updated_kwargs["scale"] = int(scale)
+                updated_kwargs["input_scale"] = int(input_scale)
+
+                replace_node.kwargs = updated_kwargs
+
+                node.replace_all_uses_with(replace_node)
+                graph.erase_node(node)
+
+                modified_graph = True
+        if modified_graph:
+            graph_module.recompile()
+            graph_module = super().call(graph_module).graph_module
+
+        return PassResult(graph_module, modified_graph)

backends/arm/_passes/convert_int64_const_ops_to_int32.py

@@ -0,0 +1,74 @@
+# Copyright 2025 Arm Limited and/or its affiliates.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-unsafe
+
+
+import logging
+
+import torch
+from executorch.backends.arm._passes.fuse_constant_ops_pass import ComputeConstantOpsAOT
+from executorch.exir.pass_base import ExportPass, PassResult
+
+
+logger = logging.getLogger(__name__)
+INT32_MIN = torch.iinfo(torch.int32).min
+INT32_MAX = torch.iinfo(torch.int32).max
+
+
+class ConvertInt64ConstOpsToInt32Pass(ExportPass):
+    """
+    Rewrite constant ops that produce int64 to int32 where safe.
+
+    List of supported operatos:
+      1. `torch.full`
+      2. `torch.arange`
+      3. `torch.eye`
+      4. `torch.linspace`
+      5. `torch.tensor`
+    """
+
+    torch_ops = [
+        torch.ops.aten.full.default,
+        torch.ops.aten.arange.default,
+        torch.ops.aten.arange.start,
+        torch.ops.aten.arange.start_step,
+        torch.ops.aten.eye.default,
+        torch.ops.aten.linspace.default,
+    ]
+
+    def call(self, graph_module: torch.fx.GraphModule):
+        modified = False
+        for node in graph_module.graph.nodes:
+            if node.op != "call_function":
+                continue
+
+            if node.target not in ComputeConstantOpsAOT.targeted_ops + self.torch_ops:
+                continue
+
+            data = node.target(*node.args, **node.kwargs)
+            if data.dtype is not torch.int64:
+                continue
+
+            min_val, max_val = torch.min(data), torch.max(data)
+            if INT32_MIN <= min_val and max_val <= INT32_MAX:
+                logger.warning(
+                    f"Casting {node.name} from torch.int64 to torch.int32"
+                    f" defined in {node.meta.get('stack_trace','[no stack trace found]')}"
+                )
+                node.update_kwarg("dtype", torch.int32)
+                modified = True
+            else:
+                logger.warning(
+                    f"[{node.name}] has values: min={min_val}, max={max_val}, which exceeds int32 range "
+                    f"([{INT32_MIN}, {INT32_MAX}]); not converting dtype to int32."
+                )
+
+        if modified:
+            graph_module.graph.eliminate_dead_code()
+            graph_module.recompile()
+            graph_module = super().call(graph_module).graph_module
+
+        return PassResult(graph_module, modified)

backends/arm/_passes/convert_int64_output_ops_to_int32.py

@@ -0,0 +1,153 @@
+# Copyright 2025 Arm Limited and/or its affiliates.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-unsafe
+
+
+import logging
+
+import torch
+from executorch.backends.arm._passes.arm_pass_utils import (
+    create_node,
+    get_first_fake_tensor,
+    set_node_arg,
+)
+from executorch.exir.dialects._ops import ops as exir_ops
+from executorch.exir.pass_base import ExportPass, PassResult
+
+
+logger = logging.getLogger(__name__)
+
+
+class ConvertInt64OutputOpsToInt32Pass(ExportPass):
+    """
+    Rewrites or removes operations that produce int64 outputs, converting them
+    to int32 where possible.
+
+
+    Currently, this pass handles casting and argmax operators:
+      1. int32 -> int64:
+         removes the cast and redirects all uses to the original int32 value.
+      2. other types -> int64:
+         rewrites the cast to produce int32 instead of int64.
+      3. torch.argmax()
+         insert an int64->int32 cast after the argmax node
+
+    Future extensions may include operators that return int64 outputs by default
+    (e.g., `argmin`), rewriting them or inserting an int64 -> int32 cast to yield
+    int32 results.
+
+    Note: Overflow checks are applied selectively in this pass. For operators without
+    such checks, it is the user's responsibility to ensure that values fit within
+    the int32 range.
+    """
+
+    aten_cast_ops = (
+        torch.ops.aten.to.dtype,
+        torch.ops.aten.to.dtype_layout,
+    )
+    edge_cast_ops = (exir_ops.edge.dim_order_ops._to_dim_order_copy.default,)
+
+    aten_argmax_ops = (torch.ops.aten.argmax.default,)
+    edge_argmax_ops = (exir_ops.edge.aten.argmax.default,)
+
+    aten_ops = aten_cast_ops + aten_argmax_ops
+    edge_ops = edge_cast_ops + edge_argmax_ops
+
+    # dtype is specified in args
+    cast_ops_args = (
+        torch.ops.aten.to.dtype,  # to_2: node.args: (gt, torch.int64) node.kwargs: {}
+    )
+    # dtype is specified in kwargs
+    cast_ops_kwargs = (
+        torch.ops.aten.to.dtype_layout,  # to_1: node.args: (unsqueeze,) node.kwargs: {'dtype': torch.int64, 'layout': torch.strided, 'device': device(type='cpu')}
+        exir_ops.edge.dim_order_ops._to_dim_order_copy.default,  # node.args: (aten_gt_scalar,) node.kwargs: {'dtype': torch.int64, 'dim_order': [0, 1]}
+    )
+
+    def _get_decomposition(self, op):
+        if op in self.edge_ops:
+            return exir_ops.edge.aten._to_copy.default
+
+        if op in self.aten_ops:
+            return torch.ops.aten._to_copy.default
+
+        raise RuntimeError(
+            f"[{self.__class__.__name__}] Can't get decomposition for op {op}"
+        )
+
+    def _convert_casting_operators(self, node: torch.fx.Node):
+        input_node = node.all_input_nodes[0]
+        input_dtype = get_first_fake_tensor(input_node).dtype
+        # Case 1: int32 -> int64 - removes the ops
+        if input_dtype == torch.int32:
+            users = [user for user in node.users if node != user]
+            for user in users:
+                logger.warning(
+                    f"Removing int32->int64 casting node {node.name} defined in"
+                    f" {node.meta.get('stack_trace','[no stack trace found]')}"
+                )
+                user.replace_input_with(node, input_node)
+        # Case 2: other types -> int64 - rewrites to cast to int32
+        else:
+            if node.target in self.cast_ops_kwargs:
+                set_node_arg(node, "dtype", torch.int32)
+            elif node.target in self.cast_ops_args:
+                set_node_arg(node, 1, torch.int32)
+            else:
+                raise RuntimeError(f"Unexpected target {node.target} in {node.name}")
+            output_dtype = get_first_fake_tensor(node).dtype
+            logger.warning(
+                f"Converting casting node {node.name} from {input_dtype}->{output_dtype} to"
+                f" {input_dtype}->torch.int32 defined in {node.meta.get('stack_trace','[no stack trace found]')}"
+            )
+
+    def _convert_argmax_operators(self, node: torch.fx.Node, graph: torch.fx.Graph):
+        output_tensor = node
+        to_copy_op = self._get_decomposition(node.target)
+        with graph.inserting_after(node):
+            cast_after = create_node(
+                graph,
+                to_copy_op,
+                args=(output_tensor,),
+                kwargs={
+                    "dtype": torch.int32,
+                },
+            )
+            users = [user for user in node.users if user != cast_after]
+            for user in users:
+                user.replace_input_with(output_tensor, cast_after)
+            logger.warning(
+                f"Inserting a casting node {cast_after.name} after {node.name} to cast int64 output"
+                f" to int32 for {node.name} defined in {node.meta.get('stack_trace','[no stack trace found]')}"
+            )
+
+    def call(self, graph_module: torch.fx.GraphModule):
+        modified = False
+        graph = graph_module.graph
+        for node in list(graph.nodes):
+            if node.op != "call_function":
+                continue
+            if node.target not in self.aten_ops + self.edge_ops:
+                continue
+            output_dtype = get_first_fake_tensor(node).dtype
+            if output_dtype != torch.int64:
+                continue
+
+            if node.target in self.aten_cast_ops + self.edge_cast_ops:
+                self._convert_casting_operators(node)
+            elif node.target in self.aten_argmax_ops + self.edge_argmax_ops:
+                # TODO: Add range check based on the input tensor shape before casting the output
+                self._convert_argmax_operators(node, graph)
+            else:
+                raise RuntimeError(f"Unexpected target {node.target} in {node.name}")
+
+            modified = True
+
+        if modified:
+            graph_module.graph.eliminate_dead_code()
+            graph_module.recompile()
+            graph_module = super().call(graph_module).graph_module
+
+        return PassResult(graph_module, modified)