Merge branch 'main' into change-1119426

Author: AdrianLundell

backends/arm/test/ops/test_stack.py

@@ -0,0 +1,150 @@
+# 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.
+
+from typing import Tuple
+
+import torch
+import torch.nn as nn
+
+from executorch.backends.arm.test import common
+from executorch.backends.arm.test.tester.test_pipeline import (
+    EthosU55PipelineINT,
+    EthosU85PipelineINT,
+    TosaPipelineFP,
+    TosaPipelineINT,
+    VgfPipeline,
+)
+
+test_data_suite = {
+    # (test_name, test_data)
+    "ones_two_tensors": lambda: ((torch.ones(1), torch.ones(1)), 0),
+    "ones_and_rand_three_tensors": lambda: (
+        (torch.ones(1, 2), torch.randn(1, 2), torch.randn(1, 2)),
+        1,
+    ),
+    "ones_and_rand_four_tensors": lambda: (
+        (
+            torch.ones(1, 2, 5),
+            torch.randn(1, 2, 5),
+            torch.randn(1, 2, 5),
+            torch.randn(1, 2, 5),
+        ),
+        -1,
+    ),
+    "rand_two_tensors": lambda: (
+        (torch.randn(2, 2, 4), torch.randn(2, 2, 4)),
+        2,
+    ),
+    "rand_two_tensors_dim_0": lambda: (
+        (torch.randn(1, 2, 4, 4), torch.randn(1, 2, 4, 4)),
+    ),
+    "rand_two_tensors_dim_2": lambda: (
+        (torch.randn(2, 2, 3, 5), torch.randn(2, 2, 3, 5)),
+        2,
+    ),
+    "rand_large": lambda: (
+        (
+            10000 * torch.randn(2, 3, 1, 4),
+            torch.randn(2, 3, 1, 4),
+            torch.randn(2, 3, 1, 4),
+        ),
+        -3,
+    ),
+}
+
+
+class Stack(nn.Module):
+    aten_op = "torch.ops.aten.stack.default"
+    exir_op = "executorch_exir_dialects_edge__ops_aten_cat_default"
+
+    def forward(self, n: tuple[torch.Tensor, ...], dim: int = 0):
+        return torch.stack(n, dim)
+
+
+input_t1 = Tuple[torch.Tensor]
+
+
+@common.parametrize("test_module", test_data_suite)
+def test_stack_tosa_FP(test_module: input_t1):
+    test_data = test_module()
+    pipeline = TosaPipelineFP[input_t1](
+        Stack(),
+        test_data,
+        aten_op=Stack.aten_op,
+        exir_op=Stack.exir_op,
+        use_to_edge_transform_and_lower=False,
+    )
+    pipeline.run()
+
+
+@common.parametrize("test_module", test_data_suite)
+def test_stack_tosa_INT(test_module: input_t1):
+    test_data = test_module()
+    pipeline = TosaPipelineINT[input_t1](
+        Stack(),
+        test_data,
+        aten_op=Stack.aten_op,
+        exir_op=Stack.exir_op,
+        use_to_edge_transform_and_lower=False,
+    )
+    pipeline.run()
+
+
+@common.XfailIfNoCorstone300
+@common.parametrize("test_module", test_data_suite)
+def test_stack_u55_INT(test_module: input_t1):
+    test_data = test_module()
+    pipeline = EthosU55PipelineINT[input_t1](
+        Stack(),
+        test_data,
+        aten_ops=Stack.aten_op,
+        exir_ops=Stack.exir_op,
+        use_to_edge_transform_and_lower=False,
+    )
+    pipeline.run()
+
+
+@common.XfailIfNoCorstone320
+@common.parametrize("test_module", test_data_suite)
+def test_stack_u85_INT(test_module: input_t1):
+    test_data = test_module()
+    pipeline = EthosU85PipelineINT[input_t1](
+        Stack(),
+        test_data,
+        aten_ops=Stack.aten_op,
+        exir_ops=Stack.exir_op,
+        use_to_edge_transform_and_lower=False,
+    )
+    pipeline.run()
+
+
+@common.SkipIfNoModelConverter
+@common.parametrize("test_module", test_data_suite)
+def test_stack_vgf_FP(test_module: input_t1):
+    test_data = test_module()
+    pipeline = VgfPipeline[input_t1](
+        Stack(),
+        test_data,
+        aten_op=Stack.aten_op,
+        exir_op=Stack.exir_op,
+        tosa_version="TOSA-1.0+FP",
+        use_to_edge_transform_and_lower=False,
+    )
+    pipeline.run()
+
+
+@common.SkipIfNoModelConverter
+@common.parametrize("test_module", test_data_suite)
+def test_stack_vgf_INT(test_module: input_t1):
+    test_data = test_module()
+    pipeline = VgfPipeline[input_t1](
+        Stack(),
+        test_data,
+        aten_op=Stack.aten_op,
+        exir_op=Stack.exir_op,
+        tosa_version="TOSA-1.0+INT",
+        use_to_edge_transform_and_lower=False,
+    )
+    pipeline.run()

backends/nxp/backend/edge_program_converter.py

@@ -18,6 +18,7 @@
 from torch.fx import Node
 from torch.nn.parameter import Parameter
 from executorch.backends.nxp.backend.ir.converter.node_converters.ops_converters import *  # noqa F403
+from executorch.backends.nxp.backend.neutron_target_spec import NeutronTargetSpec
 from executorch.backends.nxp.backend.node_format_inference import (
     NodeFormat,
     NodeFormatInference,
@@ -54,19 +55,22 @@ class EdgeProgramToIRConverter:
     """
 
     _default_conversion_config = ConversionConfig()
+    _default_target_spec = NeutronTargetSpec("imxrt700", "SDK_25_09")
     _default_delegation_options = CustomDelegationOptions()
 
     def convert_program(
         self,
         edge_program: ExportedProgram,
-        conversion_config=_default_conversion_config,
+        conversion_config: ConversionConfig = _default_conversion_config,
+        neutron_target_spec: NeutronTargetSpec = _default_target_spec,
         custom_delegation_options: CustomDelegationOptions = _default_delegation_options,
     ) -> (bytes, dict):
         """
         Convert ExportedProgram in Edge dialect to IR (TFLite flatbuffers) as bytes.
 
         :param edge_program: Converter ExportedProgram.
         :param conversion_config: ConversionConfig instance.
+        :param neutron_target_spec: Object for querying the target platform to retrieve its properties.
         :param custom_delegation_options: Custom user options which affect node delegation.
         :return: TFLite flatbuffers as bytes.
         """
@@ -76,6 +80,7 @@ def convert_program(
         cc = self.build_conversion_context(
             parameters_mapping,
             node_formats,
+            neutron_target_spec,
             conversion_config,
             custom_delegation_options,
         )
@@ -173,11 +178,12 @@ def map_inputs_to_parameters(edge_program: ExportedProgram) -> dict[str, Paramet
     def build_conversion_context(
         parameters_mapping: dict,
         node_formats: dict[Node, NodeFormat],
+        neutron_target_spec: NeutronTargetSpec,
         conversion_config: ConversionConfig = _default_conversion_config,
         custom_delegation_options: CustomDelegationOptions = _default_delegation_options,
     ) -> ConversionContext:
         tflite_builder = AtenModelBuilderDirector(
-            3, "TFLite from EdgeProgram", conversion_config
+            3, "TFLite from EdgeProgram", neutron_target_spec, conversion_config
         )
 
         # Add "sentinel" buffer (defined in schema.fbs)

backends/nxp/backend/ir/converter/builder/model_builder.py

@@ -48,6 +48,7 @@
     FlexTranspose,
 )
 from executorch.backends.nxp.backend.ir.tflite_optimizer import optimizer
+from executorch.backends.nxp.backend.neutron_target_spec import NeutronTargetSpec
 
 
 class ModelBuilder:
@@ -74,17 +75,21 @@ class ModelBuilder:
 
     _zeros_tensor_map: Dict  # Mapping 'string' shapes to 'tflT.Tensor' objects
 
-    _default_conversion_config = ConversionConfig()
+    neutron_target_spec: NeutronTargetSpec
 
     conversion_config: ConversionConfig
 
+    _default_conversion_config = ConversionConfig()
+
     def __init__(
         self,
         model_version: int,
         model_description: str,
+        neutron_target_spec: NeutronTargetSpec,
         conversion_config: ConversionConfig = _default_conversion_config,
     ) -> None:
         self._tfl_model = tflite_model.Model(model_version, model_description)
+        self.neutron_target_spec = neutron_target_spec
         self.conversion_config = conversion_config
 
         self.op_code_type_index_map = {}
@@ -471,31 +476,7 @@ def finish(self) -> tflite_model.Model:
 
         return self._tfl_model
 
-    def _assign_tensor_and_buffer_indices(  # noqa C901
-        self, allow_inputs_stripping: bool
-    ):
-        """Correctly initialize all references via indices in all tensors and buffers."""
-
-        # Assign each buffer its index
-        for i, buffer in enumerate(self.get_buffers().vector):
-            buffer.tmp_index = i
-
-        # Assign each tensor its index and its buffer index
-        for i, tensor in enumerate(self.get_tensors().vector):
-            if tensor.tmp_null_tensor:
-                # Using -1 as the index to the 'tensors' vector is way of telling the TFLite inference engine, that
-                #  this tensor should not be used.
-                # https://github.com/tensorflow/tensorflow/blob/05404d959119d41a8ffb8a75c6f232cfd8540d45/tensorflow/lite/kernels/kernel_util.cc#L79-L98
-                tensor.tmp_index = -1
-            else:
-                tensor.tmp_index = i
-
-            tensor.buffer = tensor.tmp_buffer.tmp_index
-
-        # TODO Remove inputs and outputs that are not in the tensors collection
-
-        # Assign 'Outputs' and 'Inputs' their tensor indices
-        outputs = self.get_sub_graph().outputs
+    def _assign_io_tensor_indices(self, inputs, outputs, allow_inputs_stripping: bool):
         for tensor in outputs.tmp_outputs:
             try:
                 outputs.append(tensor.tmp_index)
@@ -505,7 +486,6 @@ def _assign_tensor_and_buffer_indices(  # noqa C901
                     f"The tensor '{tensor.name}' is among the model outputs, but does NOT appear in the graph!",
                 )
 
-        inputs = self.get_sub_graph().inputs
         for tensor in inputs.tmp_inputs:
             try:
                 inputs.append(tensor.tmp_index)
@@ -520,14 +500,46 @@ def _assign_tensor_and_buffer_indices(  # noqa C901
                         f"The tensor '{tensor.name}' is among the model inputs, but does NOT appear in the graph!",
                     )
 
-        # Assign each operator its inputs and outputs indices
-        for operator in self.get_sub_graph().operators.vector:
+    def _assign_operators_io_tensor_indices(self, operators):
+        for operator in operators.vector:
             for inputTensor in operator.tmp_inputs:
                 operator.inputs.append(inputTensor.tmp_index)
 
             for outputTensor in operator.tmp_outputs:
                 operator.outputs.append(outputTensor.tmp_index)
 
+    def _assign_tensor_and_buffer_indices(self, allow_inputs_stripping: bool):
+        """Correctly initialize all references via indices in all tensors and buffers."""
+
+        # Assign each buffer its index
+        for i, buffer in enumerate(self.get_buffers().vector):
+            buffer.tmp_index = i
+
+        # Assign each tensor its index and its buffer index
+        for i, tensor in enumerate(self.get_tensors().vector):
+            if tensor.tmp_null_tensor:
+                # Using -1 as the index to the 'tensors' vector is way of telling the TFLite inference engine, that
+                #  this tensor should not be used.
+                # https://github.com/tensorflow/tensorflow/blob/05404d959119d41a8ffb8a75c6f232cfd8540d45/tensorflow/lite/kernels/kernel_util.cc#L79-L98
+                tensor.tmp_index = -1
+            else:
+                tensor.tmp_index = i
+
+            tensor.buffer = tensor.tmp_buffer.tmp_index
+
+        # TODO Remove inputs and outputs that are not in the tensors collection
+
+        subgraph = self.get_sub_graph()
+
+        # Assign 'Outputs' and 'Inputs' their tensor indices
+        self._assign_io_tensor_indices(
+            inputs=subgraph.inputs,
+            outputs=subgraph.outputs,
+            allow_inputs_stripping=allow_inputs_stripping,
+        )
+        # Assign each operator its inputs and outputs indices
+        self._assign_operators_io_tensor_indices(operators=subgraph.operators)
+
     def _build_operator_code(
         self, op_type: BuiltinOperator, version, custom_code: str = None
     ):