NXP backend: Enable permute_copy with new Neutron MLIR flow. (#19974)

### Summary
This PR updates the support for the `permute_copy` operator in the NXP
backend to reflect the requirements of the new Neutron MLIR flow. In
short, Neutron now supports all possible permutations without any
restrictions.

### Test plan
Unit tests provided.


cc @robert-kalmar @JakeStevens @digantdesai @rascani

Author: MartinPavella

backends/nxp/backend/edge_helper.py

@@ -8,27 +8,45 @@
 
 import torch
 
-from executorch.exir.dialects._ops import ops as exir_ops
+from executorch.backends.nxp.tests.ops_aliases import (
+    AddTensor,
+    Cat,
+    Clone,
+    CloneDimOrder,
+    DequantizePerChannel,
+    DequantizePerTensor,
+    MulTensor,
+    PermuteCopy,
+    QuantizePerChannel,
+    QuantizePerTensor,
+    SubTensor,
+    ViewCopy,
+)
 from torch.fx import GraphModule, Node
 from torch.fx.node import Argument
 from torch.nn import Parameter
 
 QUANTIZE_OPERATORS = [
-    exir_ops.edge.quantized_decomposed.quantize_per_channel.default,
-    exir_ops.edge.quantized_decomposed.quantize_per_tensor.default,
+    QuantizePerChannel,
+    QuantizePerTensor,
+    torch.ops.quantized_decomposed.quantize_per_tensor.default,
+    torch.ops.quantized_decomposed.quantize_per_channel.default,
 ]
 
 DEQUANTIZE_OPERATORS = [
-    exir_ops.edge.quantized_decomposed.dequantize_per_channel.default,
-    exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default,
+    DequantizePerChannel,
+    DequantizePerTensor,
+    torch.ops.quantized_decomposed.dequantize_per_tensor.default,
+    torch.ops.quantized_decomposed.dequantize_per_channel.default,
 ]
 
 # A set of operators which could possibly be no-ops in certain conditions. The operators in this set will be proclaimed
 #  as no-ops (and potentially not delegated), if their input and output tensors are equal (when run on random data).
 no_op_candidates = {
-    exir_ops.edge.aten.add.Tensor,
-    exir_ops.edge.aten.mul.Tensor,
-    exir_ops.edge.aten.sub.Tensor,
+    AddTensor,
+    MulTensor,
+    PermuteCopy,
+    SubTensor,
 }
 
 
@@ -108,21 +126,11 @@ def try_get_tensor_constant_from_node(
 
 
 def _is_dequantize(node_: Node) -> bool:
-    return node_.op == "call_function" and node_.target in [
-        exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default,
-        exir_ops.edge.quantized_decomposed.dequantize_per_channel.default,
-        torch.ops.quantized_decomposed.dequantize_per_tensor.default,
-        torch.ops.quantized_decomposed.dequantize_per_channel.default,
-    ]
+    return node_.op == "call_function" and node_.target in DEQUANTIZE_OPERATORS
 
 
 def _is_quantize(node_: Node) -> bool:
-    return node_.op == "call_function" and node_.target in [
-        exir_ops.edge.quantized_decomposed.quantize_per_tensor.default,
-        exir_ops.edge.quantized_decomposed.quantize_per_channel.default,
-        torch.ops.quantized_decomposed.quantize_per_tensor.default,
-        torch.ops.quantized_decomposed.quantize_per_channel.default,
-    ]
+    return node_.op == "call_function" and node_.target in QUANTIZE_OPERATORS
 
 
 def previous_non_qdq_node(node: Node, input_index: int = 0) -> Node | None:
@@ -172,21 +180,11 @@ def get_non_qdq_users(node: Node) -> list[Node]:
     """
 
     quant_nodes = list(node.users)
-    if len(quant_nodes) != 1 or quant_nodes[0].target not in [
-        exir_ops.edge.quantized_decomposed.quantize_per_tensor.default,
-        exir_ops.edge.quantized_decomposed.quantize_per_channel.default,
-    ]:
+    if len(quant_nodes) != 1 or not _is_quantize(quant_nodes[0]):
         return []
 
     dequant_nodes = list(quant_nodes[0].users)
-    if any(
-        dequant_node.target
-        not in [
-            exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default,
-            exir_ops.edge.quantized_decomposed.dequantize_per_channel.default,
-        ]
-        for dequant_node in dequant_nodes
-    ):
+    if any(not _is_dequantize(dequant_node) for dequant_node in dequant_nodes):
         return []
 
     res = []
@@ -277,14 +275,14 @@ def is_no_op_on_neutron(node: Node, parameters_mapping: dict[str, Parameter]) ->
         )
 
     if node.target in [
-        exir_ops.edge.aten.view_copy.default,
-        exir_ops.edge.dim_order_ops._clone_dim_order.default,
-        exir_ops.edge.aten.clone.default,
+        Clone,
+        ViewCopy,
+        CloneDimOrder,
     ]:
         # Known operators which are always no-ops on Neutron.
         return True
 
-    if node.target == exir_ops.edge.aten.cat.default and len(node.args[0]) == 1:
+    if node.target == Cat and len(node.args[0]) == 1:
         # Concatenation with 1 input is a no-op.
         return True
 

backends/nxp/backend/edge_program_converter.py

@@ -17,7 +17,7 @@
 )
 from torch._subclasses import FakeTensor
 from torch.export import ExportedProgram
-from torch.export.graph_signature import InputKind
+from torch.export.graph_signature import ExportGraphSignature, InputKind
 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
@@ -78,7 +78,7 @@ def convert_program(
         conversion_config: ConversionConfig = _default_conversion_config,
         neutron_target_spec: NeutronTargetSpec = _default_target_spec,
         custom_delegation_options: CustomDelegationOptions = _default_delegation_options,
-    ) -> tuple[bytes, dict[str, DataFormat]]:
+    ) -> tuple[bytes, dict[str, dict[str, DataFormat]]]:
         """
         Convert ExportedProgram in Edge dialect to IR (TFLite flatbuffers) as bytes.
 
@@ -95,6 +95,7 @@ def convert_program(
             parameters_mapping,
             dim_order_map,
             neutron_target_spec,
+            edge_program.graph_signature,
             conversion_config,
             custom_delegation_options,
         )
@@ -247,8 +248,9 @@ def map_nodes_to_dim_order(edge_program: ExportedProgram) -> dict[str, Parameter
     @staticmethod
     def build_conversion_context(
         parameters_mapping: dict,
-        dim_order_map: dict[str, ...],
+        dim_order_map: dict[str, Parameter],
         neutron_target_spec: NeutronTargetSpec,
+        edge_program_signature: ExportGraphSignature,
         conversion_config: ConversionConfig = _default_conversion_config,
         custom_delegation_options: CustomDelegationOptions = _default_delegation_options,
     ) -> ConversionContext:
@@ -268,6 +270,7 @@ def build_conversion_context(
             conversion_config,
             parameters_mapping,
             custom_delegation_options,
+            edge_program_signature,
         )
 
         return context

backends/nxp/backend/ir/conversion_context.py

@@ -2,14 +2,14 @@
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
-
 from executorch.backends.nxp.backend.custom_delegation_options import (
     CustomDelegationOptions,
 )
 from executorch.backends.nxp.backend.ir.conversion_config import ConversionConfig
 from executorch.backends.nxp.backend.ir.converter.builder.aten_model_builder_director import (
     AtenModelBuilderDirector,
 )
+from torch.export import ExportGraphSignature
 from torch.nn import Parameter
 
 
@@ -23,16 +23,21 @@ def __init__(
         self,
         tflite_builder: AtenModelBuilderDirector,
         conversion_config: ConversionConfig,
-        parameters_mapping: dict,
+        parameters_mapping: dict[str, Parameter],
         custom_delegation_options: CustomDelegationOptions,
+        edge_program_signature: ExportGraphSignature,
     ):
         """
         Context with data related to current conversion.
 
         :param tflite_builder: TFLite model builder.
         :param conversion_config: Conversion configuration flags and metadata.
+        :param parameters_mapping: Dictionary mapping node names to their data.
+        :param custom_delegation_options: Options that affect which nodes will be delegated.
+        :param edge_program_signature: Description of the inputs of the edge graph.
         """
         self.tflite_builder = tflite_builder
         self.conversion_config = conversion_config
         self.parameters_mapping = parameters_mapping
         self.custom_delegation_options = custom_delegation_options
+        self.edge_program_signature = edge_program_signature

backends/nxp/backend/ir/converter/node_converters/ops_converters/getitem_converter.py

@@ -41,5 +41,16 @@ def convert(self, node: Node):
             input_.type = output.type
             propagate_quantization(from_tensor=output, to_tensor=input_)
 
-        self.builder.turn_operator_to_identity(t_op)
-        self.builder.append_operators([t_op])
+        consumes_model_input = (
+            node.args[0].name in self.context.edge_program_signature.user_inputs
+        )
+        if consumes_model_input:
+            # Convert as identity op (Transpose that will be removed) because the input tensor is also an input of the
+            #  model. If we did redirection here, we would change the name of a model input, which is prohibited.
+            self.builder.turn_operator_to_identity(t_op)
+            self.builder.append_operators([t_op])
+        else:
+            # The operator will be converted to nothing. That means its output will not be in the model. We need to
+            #  redirect the output to the input, so that any operators that consume the `output` will use the `input_`
+            #  instead.
+            self.builder.redirect_tensor(output, input_)