Update on "add more export modules after ertrecod created"

we need to support etrecord recording custom export modules for further usage. This diff makes that happen by creating new function inside ETRecord

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

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