Update on "[executorch][core] NamedDataMap interface"

Add NamedDataMap interface to runtime.

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

[ghstack-poisoned]

Author: lucylq

.ci/docker/requirements-ci.txt

@@ -1,5 +1,5 @@
 mpmath==1.3.0
-numpy==2.0.0; python_version >= '3.10'
+numpy>=2.0.0; python_version >= '3.10'
 PyYAML==6.0.1
 ruamel.yaml==0.17.32
 sympy==1.12
@@ -8,7 +8,7 @@ tomli==2.0.1
 torchsr==1.0.4
 transformers==4.47.1
 zstd==1.5.5.1
-pandas==2.2.2; python_version >= '3.10'
+pandas>=2.2.2; python_version >= '3.10'
 pytest==7.2.0
 pytest-cov==4.1.0
 expecttest==0.1.6
@@ -21,7 +21,7 @@ sphinx-gallery==0.14.0
 breathe==4.34.0
 exhale==0.2.3
 docutils==0.16
-matplotlib==3.9.4
+matplotlib>=3.9.4
 # PyTorch Theme
 -e git+https://github.com/pytorch/pytorch_sphinx_theme.git#egg=pytorch_sphinx_theme
 myst-parser==0.18.1

backends/arm/_passes/annotate_channels_last_dim_order_pass.py

@@ -12,7 +12,6 @@
 from executorch.backends.arm._passes.arm_pass_utils import (
     create_node,
     get_first_fake_tensor,
-    get_node_arg,
     insert_q_dq_pair,
 )
 from executorch.backends.arm.tosa_quant_utils import dq_op, q_op
@@ -26,9 +25,8 @@
 # when lowering to TOSA, e.g. a passthrough_to_tosa._transpose will not affect
 # the edge IR graph but will be lowered to a TOSA-TRANSPOSE.
 lib = Library("passthrough_to_tosa", "DEF")
-# For operators that change the rank of the input, such as unsqueeze and squeeze, we may need
-# to switch dim_order before the opertation. Changing tosa_dim_order is not sufficient
-# as we also need transpose the data into the correct data format.
+# For certain operators we need the data in a specific data format. Changing tosa_dim_order
+# is not sufficient as we also need transpose the data.
 # By utilizing an edge IR passthrough operator we can keep the edge program in
 # channels-first/contiguous and get the desired behavior in the TOSA lowering.
 lib.define("_transpose(Tensor self, int[] dim_order) -> Tensor")
@@ -153,27 +151,6 @@ def insert_output_transpose(node, graph_module):
                 q_params = node.args[0].args[1:]
                 insert_q_dq_pair(graph_module.graph, node, q_params)
 
-    @staticmethod
-    def _insert_squeeze_transpose(
-        input_shape, output_shape, node, input_node, graph_module
-    ):
-        nhwc_to_nhwc = len(input_shape) == 4 and len(output_shape) <= 3
-
-        if nhwc_to_nhwc and AnnotateChannelsLastDimOrder.memory_format_differs(
-            input_shape
-        ):
-            AnnotateChannelsLastDimOrder.insert_input_transpose(
-                node, input_node, graph_module
-            )
-
-    @staticmethod
-    def _insert_unsqueeze_transpose(input_shape, output_shape, node, graph_module):
-        nchw_to_nhwc = len(input_shape) == 3 and len(output_shape) == 4
-        if nchw_to_nhwc and AnnotateChannelsLastDimOrder.memory_format_differs(
-            output_shape
-        ):
-            AnnotateChannelsLastDimOrder.insert_output_transpose(node, graph_module)
-
     @staticmethod
     def _insert_view_transpose(
         input_shape, output_shape, node, input_node, graph_module
@@ -199,8 +176,6 @@ def insert_tosa_transposes(self, graph_module: torch.fx.GraphModule):
         """
         Transposes are needed for operators transforming the input to a different rank, as 4D-tensors are assumed to be in NHWC-format, whereas all other are in NCHW format.
         This is relevant for the following cases:
-        - squeeze:     4D -> <4D
-        - unsqueeze:   3D ->  4D
         - view:       <4D ->  4D
         - view:        4D -> <4D
         Additionally, a 4D->4D view operation acting on the channel dimension currently needs to be performed in NCHW format, leadning to one extra input and output transpose for this case.
@@ -214,27 +189,6 @@ def insert_tosa_transposes(self, graph_module: torch.fx.GraphModule):
             if node.op != "call_function":
                 continue
 
-            if node.target == exir_ops.edge.aten.squeeze_copy.dims:
-                input_node = node.args[0]
-                input_shape = input_node.meta["val"].shape
-                output_shape = node.meta["val"].shape
-
-                self._insert_squeeze_transpose(
-                    input_shape, output_shape, node, input_node, graph_module
-                )
-
-            elif node.target == exir_ops.edge.aten.unsqueeze_copy.default:
-                input_node = get_node_arg(node.args, 0, default_value=False)
-                if input_node:
-                    input_shape = input_node.meta["val"].shape
-                else:
-                    input_shape = ()
-                output_shape = node.meta["val"].shape
-
-                self._insert_unsqueeze_transpose(
-                    input_shape, output_shape, node, graph_module
-                )
-
             elif node.target == exir_ops.edge.aten.view_copy.default:
                 input_node = node.args[0]
                 input_shape = input_node.meta["val"].shape

backends/arm/_passes/arm_pass_manager.py

@@ -21,6 +21,9 @@
 from executorch.backends.arm._passes.convert_split_to_slice import (
     ConvertSplitToSlicePass,
 )
+from executorch.backends.arm._passes.convert_squeezes_to_view import (
+    ConvertSqueezesToViewPass,
+)
 from executorch.backends.arm._passes.decompose_div_pass import DecomposeDivPass
 from executorch.backends.arm._passes.decompose_layernorm_pass import (
     DecomposeLayerNormPass,
@@ -100,6 +103,7 @@ def _tosa_080_BI_pipeline(self, exported_program: ExportedProgram) -> GraphModul
         self.add_pass(KeepDimsFalseToSqueezePass())
         self.add_pass(Conv1dUnsqueezePass(exported_program))
         self.add_pass(DecomposeSelectPass())
+        self.add_pass(ConvertSqueezesToViewPass())
 
         self.add_pass(AnnotateChannelsLastDimOrder())
 
@@ -135,6 +139,7 @@ def _tosa_080_MI_pipeline(self, exported_program: ExportedProgram) -> GraphModul
         self.add_pass(KeepDimsFalseToSqueezePass())
         self.add_pass(Conv1dUnsqueezePass(exported_program))
         self.add_pass(DecomposeSelectPass())
+        self.add_pass(ConvertSqueezesToViewPass())
 
         self.add_pass(AnnotateChannelsLastDimOrder())
 

backends/arm/_passes/convert_squeezes_to_view.py

@@ -0,0 +1,30 @@
+# Copyright 2025 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.
+
+# pyre-unsafe
+
+from executorch.exir.dialects._ops import ops as exir_ops
+from executorch.exir.pass_base import ExportPass
+
+
+class ConvertSqueezesToViewPass(ExportPass):
+    """
+    Replaces squeeze/unsqueeze operators with view. These are simply special cases of the view op, so removing them gives us less cases to handle in the node visitiors.
+    """
+
+    def call_operator(self, op, args, kwargs, meta):
+        if op not in [
+            exir_ops.edge.aten.squeeze_copy.dims,
+            exir_ops.edge.aten.unsqueeze_copy.default,
+        ]:
+            return super().call_operator(op, args, kwargs, meta)
+
+        x = args[0]
+        shape = meta["val"].size()
+        view_args = (x, list(shape))
+        return super().call_operator(
+            exir_ops.edge.aten.view_copy.default, view_args, kwargs, meta
+        )

backends/arm/operator_support/to_copy_support.py

@@ -1,4 +1,4 @@
-# Copyright 2024 Arm Limited and/or its affiliates.
+# Copyright 2024-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.
@@ -22,7 +22,10 @@
 
 @register_tosa_support_check
 class ToCopySupported(SupportedTOSAOperatorCheck):
-    targets = [exir_ops.edge.aten._to_copy.default]
+    targets = [
+        exir_ops.edge.aten._to_copy.default,
+        exir_ops.edge.dim_order_ops._to_dim_order_copy.default,
+    ]
 
     tosa_specs = [
         TosaSpecification.create_from_string("TOSA-0.80+BI"),
@@ -110,7 +113,7 @@ def is_node_supported(self, node: fx.Node, tosa_spec: TosaSpecification) -> bool
             )
             return False
 
-        # Check memory format
+        # Check memory format (to_copy)
         if "memory_format" in node.kwargs:
             if node.kwargs["memory_format"] in (torch.preserve_format,):
                 logger.info(
@@ -119,4 +122,14 @@ def is_node_supported(self, node: fx.Node, tosa_spec: TosaSpecification) -> bool
                 )
                 return False
 
+        # Check dim_order (to_dim_order_copy)
+        if "dim_order" in node.kwargs:
+            dim_order = node.kwargs["dim_order"]
+            if dim_order != list(range(len(dim_order))):
+                logger.info(
+                    f"Argument {dim_order=} is not supported for "
+                    f"{node.target.name()} right now."  # pyre-ignore[16]
+                )
+                return False
+
         return True

backends/arm/operator_support/tosa_supported_operators.py

@@ -1,4 +1,4 @@
-# Copyright 2024 Arm Limited and/or its affiliates.
+# Copyright 2024-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.
@@ -81,11 +81,16 @@ def is_node_supported(self, submodules, node: fx.Node) -> bool:
             exir_ops.edge.aten.hardtanh.default,
             exir_ops.edge.aten.convolution.default,
             exir_ops.edge.aten.div.Tensor,
+            exir_ops.edge.aten.eq.Tensor,
             exir_ops.edge.aten.exp.default,
             exir_ops.edge.aten.log.default,
             exir_ops.edge.aten.linear.default,
             exir_ops.edge.aten.split_with_sizes_copy.default,
             exir_ops.edge.aten.full.default,
+            exir_ops.edge.aten.ge.Tensor,
+            exir_ops.edge.aten.gt.Tensor,
+            exir_ops.edge.aten.le.Tensor,
+            exir_ops.edge.aten.lt.Tensor,
             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,

backends/arm/operators/__init__.py

@@ -13,11 +13,16 @@
     op_bmm,
     op_cat,
     op_conv2d,
+    op_eq,
     op_exp,
     op_full,
+    op_ge,
     op_get_item,
+    op_gt,
     op_hardtanh,
+    op_le,
     op_log,
+    op_lt,
     op_max,
     op_max_pool2d,
     op_min,
@@ -30,14 +35,13 @@
     op_rsqrt,
     op_sigmoid,
     op_slice,
-    op_squeeze,
     op_sub,
     op_sum,
     op_table,
     op_tanh,
     op_to_copy,
+    op_to_dim_order_copy,
     op_transpose,
-    op_unsqueeze,
     op_upsample_nearest2d,
     op_view,
 )

backends/arm/operators/op_eq.py

@@ -0,0 +1,58 @@
+# 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 typing import List
+
+import executorch.backends.arm.tosa_quant_utils as tqutils
+
+import serializer.tosa_serializer as ts
+from executorch.backends.arm.operators.node_visitor import (
+    NodeVisitor,
+    register_node_visitor,
+)
+from executorch.backends.arm.tosa_mapping import TosaArg
+from serializer.tosa_serializer import TosaOp
+
+from torch.fx import Node
+
+
+@register_node_visitor
+class EqualVisitor(NodeVisitor):
+    target = "aten.eq.Tensor"
+
+    def __init__(self, *args):
+        super().__init__(*args)
+
+    def define_node(
+        self,
+        node: Node,
+        tosa_graph: ts.TosaSerializer,
+        inputs: List[TosaArg],
+        output: TosaArg,
+    ) -> None:
+        assert (
+            inputs[0].dtype == inputs[1].dtype
+        ), "EQ must have the same dtypes as input"
+
+        input_nodes = inputs
+        # Handle quantization
+        if inputs[0].dtype == ts.DType.INT8:
+            # Rescale inputs to 32 bit
+            rescaled_inputs, _ = tqutils.insert_rescale_ops_to_int32(
+                tosa_graph, inputs, node
+            )
+
+            # Update IO
+            input_nodes = rescaled_inputs
+
+        # Do the equal comparison
+        tosa_graph.addOperator(
+            TosaOp.Op().EQUAL,
+            [input_nodes[0].name, input_nodes[1].name],
+            output.name,
+            None,
+        )

backends/arm/operators/op_ge.py

@@ -0,0 +1,57 @@
+# 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 typing import List
+
+import executorch.backends.arm.tosa_quant_utils as tqutils
+
+import serializer.tosa_serializer as ts
+from executorch.backends.arm.operators.node_visitor import (
+    NodeVisitor,
+    register_node_visitor,
+)
+from executorch.backends.arm.tosa_mapping import TosaArg
+from serializer.tosa_serializer import TosaOp
+
+from torch.fx import Node
+
+
+@register_node_visitor
+class GreaterEqualVisitor(NodeVisitor):
+    target = "aten.ge.Tensor"
+
+    def __init__(self, *args):
+        super().__init__(*args)
+
+    def define_node(
+        self,
+        node: Node,
+        tosa_graph: ts.TosaSerializer,
+        inputs: List[TosaArg],
+        output: TosaArg,
+    ) -> None:
+        assert (
+            inputs[0].dtype == inputs[1].dtype
+        ), "GE must have the same dtypes as input"
+
+        input_nodes = inputs
+        # Handle quantization
+        if inputs[0].dtype == ts.DType.INT8:
+            # Rescale inputs to 32 bit
+            rescaled_inputs, _ = tqutils.insert_rescale_ops_to_int32(
+                tosa_graph, inputs, node
+            )
+
+            # Update IO
+            input_nodes = rescaled_inputs
+
+        tosa_graph.addOperator(
+            TosaOp.Op().GREATER_EQUAL,
+            [input_nodes[0].name, input_nodes[1].name],
+            [output.name],
+            None,
+        )