Update on "make to_edge_transform_and_lower support etrecord generation"

Differential Revision: [D79336982](https://our.internmc.facebook.com/intern/diff/D79336982/)

umbrella issue: #12961

[ghstack-poisoned]

Author: Gasoonjia

backends/arm/README.md

@@ -181,8 +181,8 @@ The Arm EthosU Backend should be considered a prototype quality at this point, l
 ## Current flows
 
 The EthosUBackend has a two stage process,
-- Compile to TOSA to rationalise the graph into known hardware support profiles. Currently this is to v0.80 TOSA BI with specific concern to a subset which gives support on Ethos-U55 and Ethos-U85, the target of the initial prototype efforts. This calls into the TOSABackend.
-- Lower via the ethos-u-vela compilation flow which takes TOSA v0.80 as an input and produces a low level commandstream for the hardware which is then passed via the delegate to the ethos-u-core-driver for direct execution.
+- Compile to TOSA to rationalise the graph into known hardware support profiles. Currently this is to v1.0 TOSA INT with specific concern to a subset which gives support on Ethos-U55 and Ethos-U85, the target of the initial prototype efforts. This calls into the TOSABackend.
+- Lower via the ethos-u-vela compilation flow which takes TOSA v1.0 as an input and produces a low level commandstream for the hardware which is then passed via the delegate to the ethos-u-core-driver for direct execution.
 
 The EthosUPartitioner is currenly used to ensure the operations converted are Ethos-U compatible, but will be extended to offer spec-correct TOSA Base inference and TOSA Main Inference generation in future.
 

backends/arm/TARGETS

@@ -1,10 +1,20 @@
 # @noautodeps
 load("@fbcode_macros//build_defs:python_library.bzl", "python_library")
+
+python_library(
+    name = "ethosu_partitioner",
+    srcs = [
+        "ethosu/__init__.py",
+        "ethosu/backend.py",
+        "ethosu/partitioner.py"
+    ],
+    deps = [
+        ":arm_partitioner",
+    ]
+)
 python_library(
     name = "arm_partitioner",
     srcs = [
-        "ethosu_backend.py",
-        "ethosu_partitioner.py",
         "tosa_backend.py",
         "tosa_partitioner.py",
         "vgf_backend.py",

backends/arm/_passes/__init__.py

@@ -26,6 +26,7 @@
 from .decompose_adaptive_avg_pool2d_pass import DecomposeAdaptiveAvgPool2dPass  # noqa
 from .decompose_addmm_pass import DecomposeAddmmPass  # noqa
 from .decompose_asin_pass import DecomposeAsinPass  # noqa
+from .decompose_asinh_pass import DecomposeAsinhPass  # noqa
 from .decompose_atan_pass import DecomposeAtanPass  # noqa
 from .decompose_atanh_pass import DecomposeAtanhPass  # noqa
 from .decompose_avg_pool2d import DecomposeAvgPool2d  # noqa

backends/arm/_passes/arm_pass_manager.py

@@ -30,6 +30,7 @@
     DecomposeAcoshPass,
     DecomposeAdaptiveAvgPool2dPass,
     DecomposeAddmmPass,
+    DecomposeAsinhPass,
     DecomposeAsinPass,
     DecomposeAtanhPass,
     DecomposeAtanPass,
@@ -105,7 +106,7 @@ def _transform(self, graph_module: GraphModule):
         with TosaLoweringContext(self.tosa_spec):
             return self(graph_module).graph_module
 
-    def _tosa_080_BI_pipeline(self, exported_program: ExportedProgram) -> GraphModule:
+    def _tosa_INT_pipeline(self, exported_program: ExportedProgram) -> GraphModule:
         self.add_pass(FuseQuantizedActivationPass())
         self.add_pass(RemoveGetItemPass())
         self.add_pass(ConvertSplitToSlicePass())
@@ -114,7 +115,6 @@ def _tosa_080_BI_pipeline(self, exported_program: ExportedProgram) -> GraphModul
         self.add_pass(
             DecomposeMeanDimPass(exported_program.graph_module, self.tosa_spec)
         )
-
         self.add_pass(ConvertFullLikeToFullPass())
         self.add_pass(ConvertToClampPass())
         self.add_pass(ConvertMinMaxPass())
@@ -148,7 +148,6 @@ def _tosa_080_BI_pipeline(self, exported_program: ExportedProgram) -> GraphModul
         self.add_pass(DecomposeMaxPool2DPass())
         self.add_pass(SizeAdjustInputPass())
         self.add_pass(DecomposeSelectPass())
-
         self.add_pass(ConvertSqueezesToViewPass())
 
         self.add_pass(FuseViewCopyTransform())
@@ -162,11 +161,12 @@ def _tosa_080_BI_pipeline(self, exported_program: ExportedProgram) -> GraphModul
 
         return self._transform(exported_program.graph_module)
 
-    def _tosa_080_MI_pipeline(self, exported_program: ExportedProgram) -> GraphModule:
+    def _tosa_FP_pipeline(self, exported_program: ExportedProgram) -> GraphModule:
         self.add_pass(DecomposeMaskedFill())
         self.add_pass(DecomposeRoundPass())
         self.add_pass(DecomposeAcoshPass())
         self.add_pass(DecomposeAsinPass())
+        self.add_pass(DecomposeAsinhPass())
         self.add_pass(DecomposeSqrtPass())
         self.add_pass(DecomposeAtanPass())
         self.add_pass(DecomposeAtanhPass())
@@ -235,22 +235,12 @@ def _tosa_080_MI_pipeline(self, exported_program: ExportedProgram) -> GraphModul
 
         return self._transform(exported_program.graph_module)
 
-    def _tosa_1_0_int_quantized_pipeline(self, exported_program: ExportedProgram):
-        return self._tosa_080_BI_pipeline(exported_program)
-
-    def _tosa_1_0_fp_pipeline(self, exported_program: ExportedProgram):
-        return self._tosa_080_MI_pipeline(exported_program)
-
     def transform_to_backend_pipeline(self, exported_program: ExportedProgram):
         """Apply passes before transforming program to backend"""
-        if self.tosa_spec == TosaSpecification.create_from_string("TOSA-0.80.0+BI"):
-            return self._tosa_080_BI_pipeline(exported_program)
-        elif self.tosa_spec == TosaSpecification.create_from_string("TOSA-0.80.0+MI"):
-            return self._tosa_080_MI_pipeline(exported_program)
-        elif self.tosa_spec == TosaSpecification.create_from_string("TOSA-1.0+FP"):
-            return self._tosa_1_0_fp_pipeline(exported_program)
+        if self.tosa_spec == TosaSpecification.create_from_string("TOSA-1.0+FP"):
+            return self._tosa_FP_pipeline(exported_program)
         elif self.tosa_spec == TosaSpecification.create_from_string("TOSA-1.0+INT"):
-            return self._tosa_1_0_int_quantized_pipeline(exported_program)
+            return self._tosa_INT_pipeline(exported_program)
         else:
             raise NotImplementedError(
                 f"No pass pipeline implemented for {self.tosa_spec=}"

backends/arm/_passes/decompose_asinh_pass.py

@@ -0,0 +1,50 @@
+# 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
+
+
+from executorch.backends.arm._passes import ArmPass
+from executorch.exir.dialects._ops import ops as exir_ops
+
+# For MI case
+edge_asinh_op = (exir_ops.edge.aten.asinh.default,)
+
+
+class DecomposeAsinhPass(ArmPass):
+    """
+    Decomposes asinh to supported TOSA-operations.
+    This decomposition is based on the mathematical identity:
+        asinh(x) = log(x + sqrt(x^2 + 1))
+    """
+
+    def call_operator(self, op, args, kwargs, meta):
+        if op not in edge_asinh_op:
+            return super().call_operator(op, args, kwargs, meta)
+
+        log_op, sqrt_op, mul_op, add_op_scalar, add_op = (
+            exir_ops.edge.aten.log.default,
+            exir_ops.edge.aten.sqrt.default,
+            exir_ops.edge.aten.mul.Tensor,
+            exir_ops.edge.aten.add.Scalar,
+            exir_ops.edge.aten.add.Tensor,
+        )
+
+        x = args[0]
+
+        # calculate t1 = x^2 + 1
+        x2 = super().call_operator(mul_op, (x, x), {}, meta, True)
+        t1 = super().call_operator(add_op_scalar, (x2, 1.0), {}, meta, True)
+
+        # t2 = sqrt(t1)
+        t2 = super().call_operator(sqrt_op, (t1,), {}, meta, True)
+
+        # t3 = x + t2
+        t3 = super().call_operator(add_op, (x, t2), {}, meta, True)
+
+        # out = ln(t3)
+        out = super().call_operator(log_op, (t3,), {}, meta, True)
+
+        return out

backends/arm/_passes/fuse_constant_ops_pass.py

@@ -6,6 +6,7 @@
 import logging
 
 import torch._export.utils
+import torch.fx
 from executorch.backends.arm._passes.arm_pass_utils import (
     get_constant_placeholder_kind,
     get_first_fake_tensor,
@@ -50,22 +51,26 @@ def _fuse_nodes(self, node) -> bool:
         the operations already carried out on the data.
         """
 
-        # Extract tensors and args from the node
-        data_list = [
-            get_param_tensor(self.exported_program, input_node)
-            for input_node in node.all_input_nodes
-        ]
-
-        args = node.args[len(node.all_input_nodes) :]
-        kwargs = node.kwargs
-
-        if "input_qparams" in node.meta and len(node.meta["input_qparams"]) > 0:
-            for i in range(len(node.all_input_nodes)):
-                q_params = node.meta["input_qparams"][i]
-                data_list[i] = q_params.dequantize_value(data_list[i])
-
-        # Run the op on the extracted tensor
-        data = node.target(*data_list, *args, **kwargs)
+        input_nodes = list(node.all_input_nodes)
+        qparams = node.meta.get("input_qparams", None)
+
+        def resolve_arg(arg):
+            if isinstance(arg, torch.fx.Node) and arg in input_nodes:
+                idx = input_nodes.index(arg)
+                t = get_param_tensor(self.exported_program, arg)
+                if qparams:
+                    t = qparams[idx].dequantize_value(t)
+                return t
+            if isinstance(arg, tuple):
+                return tuple(resolve_arg(x) for x in arg)
+            if isinstance(arg, list):
+                return [resolve_arg(x) for x in arg]
+            return arg
+
+        new_args = tuple(resolve_arg(a) for a in node.args)
+        new_kwargs = {k: resolve_arg(v) for k, v in node.kwargs.items()}
+
+        data = node.target(*new_args, **new_kwargs)
 
         # Only fuse if the tensor does not get bigger.
         if data.numel() > get_first_fake_tensor(node).numel():

backends/arm/_passes/insert_table_ops.py

@@ -58,6 +58,7 @@ class TableOps:
         exir_ops.edge.aten.sinh.default: torch.sinh,
         exir_ops.edge.aten.acosh.default: torch.acosh,
         exir_ops.edge.aten.asin.default: torch.asin,
+        exir_ops.edge.aten.asinh.default: torch.asinh,
     }
 
     # Targets that must be treated explicitly

backends/arm/ethosu/__init__.py

@@ -0,0 +1,14 @@
+# 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
+
+from .backend import EthosUBackend  # noqa: F401
+from .partitioner import EthosUPartitioner  # noqa: F401
+
+__all__ = [
+    "EthosUBackend",
+    "EthosUPartitioner",
+]

backends/arm/ethosu_backend.py renamed to backends/arm/ethosu/backend.py

@@ -10,7 +10,7 @@
 from executorch.backends.arm.arm_backend import (
     is_ethosu,
 )  # usort: skip
-from executorch.backends.arm.ethosu_backend import EthosUBackend
+from executorch.backends.arm.ethosu import EthosUBackend
 from executorch.backends.arm.tosa_partitioner import TOSAPartitioner
 from executorch.exir.backend.compile_spec_schema import CompileSpec
 from executorch.exir.backend.partitioner import DelegationSpec