Pull request pytorch#42: Feature/EIEX-84 improve nodeformat inference to handle format change for op aten view copy

Merge in AITEC/executorch from feature/EIEX-84-improve-nodeformat-inference-to-handle-format-change-for-op-aten-view_copy to main-nxp

* commit 'db884b4b15f4dafd930769305cc9e0d05b5856bf':
  [NO-UPSTREAM] Add tests for proper handling of node formats of `aten.view_copy` (Reshape) in TFLite.
  Update node format inference to properly support `aten.view_copy`.

Author: Pop-korn

backends/nxp/backend/node_format_inference.py

@@ -5,12 +5,10 @@
 # LICENSE file in the root directory of this source tree.
 
 import logging
-from collections import abc
 from enum import Enum
 
 from torch import Node
 from torch.export import ExportedProgram
-from typing_extensions import Tuple
 
 # (TODO Lukas) Can we found ops somewhere else?
 from executorch.exir.dialects._ops import ops as exir_ops
@@ -45,7 +43,7 @@ class NodeFormatInference:
     # A set of Edge Aten ops, which have the ability to change the format (for example - input nodes
     # are channels first but output is formatless).
     ops_that_can_change_tensor_format = {
-        # TODO ("transpose", "reshape", etc.)
+        exir_ops.edge.aten.view_copy.default
     }
 
     _node_format_mapping: dict[Node, NodeFormat]
@@ -86,13 +84,55 @@ def _infer_format_of_nodes(self, node: Node):
         if op_type in self.ops_with_channels_first_nodes:
             self._handle_node_which_uses_channels_first_format(node)
         elif op_type in self.ops_that_can_change_tensor_format:
-            if op_type in ["transpose"]:
-                self._assign_format_to_node(node, NodeFormat.FORMATLESS)
+            if op_type == exir_ops.edge.aten.view_copy.default:  # view_copy
+                self._assign_format_to_node(self._node_outputs[node][0], NodeFormat.FORMATLESS)
             else:
                 logger.error(f"Node format inference for node type: {op_type} not found!")
         else:
             self._handle_node_which_can_use_any_node_format(node)
 
+    def _infer_format_based_on_io_ranks(self, node: Node):
+        """ Determine the format of the output tensor of given "reshape style operator" based on the ranks of its input
+             and output.
+        """
+        # noinspection PyBroadException
+        try:
+            main_input_rank = len(node.all_input_nodes[0].meta['val'].shape)
+            main_output_rank = len(node.meta['val'].shape)
+
+            if main_output_rank == main_input_rank:
+                # Operator maintains the number of dimensions -> try to propagate the format.
+                self._match_formats_of_nodes(node, node.prev)
+
+            else:
+                # Either the op 'flattens' the tensor, so output is formatless, or it scales it up, in which case the
+                # format is assumed to be 'FORMATLESS', and may be back propagated as channels first later.
+                self._assign_format_to_node(node, NodeFormat.FORMATLESS)
+
+        except:
+            # Some shape data is not known, so we cannot be extra clever. Just set the output to `FORMATLESS` and
+            #  everything will be alright.
+            self._assign_format_to_node(node, NodeFormat.FORMATLESS)
+
+    def _match_formats_of_nodes(self, node_1, node_2):
+        """ If one of 'node_1' or 'node_2' is channels first, make the other channels first as well.
+             If neither is channels first, make them both formatless.
+        """
+
+        format_1 = self._get_node_format(node_1)
+        format_2 = self._get_node_format(node_2)
+
+        if format_1.is_channels_first() or format_2.is_channels_first():
+            # At least 1 is channels first
+            if not format_1.is_channels_first():
+                self._assign_format_to_node(node_1, NodeFormat.CHANNELS_FIRST)
+            elif not format_2.is_channels_first():
+                self._assign_format_to_node(node_2, NodeFormat.CHANNELS_FIRST)
+
+        else:
+            self._assign_format_to_node(node_1, NodeFormat.FORMATLESS)
+            self._assign_format_to_node(node_2, NodeFormat.FORMATLESS)
+
     def _assign_format_to_node(self, node: Node, node_format: NodeFormat):
         """
         Assign format to node, but only if it's not channels first.

backends/nxp/tests/executorch_pipeline.py

@@ -1,6 +1,8 @@
 import torch
+from torch import nn
 from torch.ao.quantization.quantize_pt2e import prepare_pt2e, convert_pt2e
 
+from executorch import exir
 from executorch.backends.nxp.neutron_partitioner import NeutronPartitioner
 from executorch.backends.nxp.nxp_backend import generate_neutron_compile_spec
 from executorch.backends.nxp.quantizer.neutron_quantizer import NeutronQuantizer
@@ -41,3 +43,9 @@ def to_quantized_executorch_program(model: torch.nn.Module, input_shape: tuple)
             extract_delegate_segments=False, extract_constant_segment=False
         )
     )
+
+
+def to_edge_program(model: nn.Module, input_shape) -> EdgeProgramManager:
+    example_input = (torch.ones(input_shape),)
+    exir_program = torch.export.export(model, example_input)
+    return exir.to_edge(exir_program)

backends/nxp/tests/ir/converter/node_converter/test_view_copy_converter.py

@@ -0,0 +1,132 @@
+from typing import Sequence
+
+import numpy as np
+import torch
+from torch import nn
+
+from executorch.backends.nxp.backend.ir.converter.builder.model_builder import ModelBuilder
+from executorch.backends.nxp.backend.ir.tflite_generator.builtin_options.conv_2d_options import Conv2D
+from executorch.backends.nxp.backend.ir.tflite_generator.builtin_options.reshape_options import Reshape
+from executorch.backends.nxp.backend.ir.tflite_generator.builtin_options.transpose_options import Transpose
+from executorch.backends.nxp.tests.executorch_pipeline import to_edge_program
+from executorch.backends.nxp.tests.executors import convert_run_compare, ToNHWCPreprocess, ToNCHWPreprocess
+
+
+class FormatlessToChannelsFirstModule(nn.Module):
+    def __init__(self, channels: int, new_shape: Sequence[int]):
+        super().__init__()
+        self.conv = nn.Conv2d(channels, channels, 2, bias=True)
+        self.new_shape = new_shape
+
+    def forward(self, x):
+        x = torch.reshape(x, self.new_shape)
+        x = self.conv(x)
+        return x
+
+
+class FormatlessToFormatlessModule(nn.Module):
+    def __init__(self, new_shape: Sequence[int]):
+        super().__init__()
+        self.new_shape = new_shape
+
+    def forward(self, x):
+        x = torch.reshape(x, self.new_shape)
+        return x
+
+
+class ConvReshapeModule(nn.Module):
+    def __init__(self, channels: int, new_shape: Sequence[int]):
+        super().__init__()
+        self.conv = nn.Conv2d(channels, channels, 2, bias=True)
+        self.new_shape = new_shape
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = torch.reshape(x, self.new_shape)
+        return x
+
+
+def test__channels_first_to_2d(mocker):
+    input_shape = [2, 4, 7, 9]
+    new_shape = [12, 32]  # Mix up the dimensions for a thorough test.
+
+    torch_model = ConvReshapeModule(channels=input_shape[1], new_shape=new_shape)
+    edge_program = to_edge_program(torch_model, input_shape).exported_program()
+
+    torch.manual_seed(23)
+    input_data = np.random.random(input_shape).astype('float32')
+
+    converter_spy = mocker.spy(ModelBuilder, "finish")
+
+    convert_run_compare(edge_program, input_data, tflite_input_preprocess=ToNHWCPreprocess())
+
+    tflite_model = converter_spy.spy_return
+    ops = tflite_model.sub_graphs[0].operators.vector
+    assert len(ops) == 3
+    assert isinstance(ops[0].builtin_options, Conv2D)
+    assert isinstance(ops[1].builtin_options, Transpose)
+    assert isinstance(ops[2].builtin_options, Reshape)
+
+
+def test__channels_first_to_4d(mocker):
+    input_shape = [2, 4, 6, 8]
+    new_shape = [7, 4, 2, 5]
+
+    torch_model = ConvReshapeModule(channels=input_shape[1], new_shape=new_shape)
+    edge_program = to_edge_program(torch_model, input_shape).exported_program()
+
+    torch.manual_seed(23)
+    input_data = np.random.random(input_shape).astype('float32')
+
+    converter_spy = mocker.spy(ModelBuilder, "finish")
+
+    convert_run_compare(edge_program, input_data, tflite_input_preprocess=ToNHWCPreprocess())
+
+    tflite_model = converter_spy.spy_return
+    ops = tflite_model.sub_graphs[0].operators.vector
+    assert len(ops) == 3
+    assert isinstance(ops[0].builtin_options, Conv2D)
+    assert isinstance(ops[1].builtin_options, Transpose)
+    assert isinstance(ops[2].builtin_options, Reshape)
+
+
+def test__formatless_to_channels_first(mocker):
+    input_shape = [12, 32]
+    new_shape = [2, 4, 6, 8]  # Mix up the dimensions for a thorough test.
+
+    torch_model = FormatlessToChannelsFirstModule(channels=new_shape[1], new_shape=new_shape)
+    edge_program = to_edge_program(torch_model, input_shape).exported_program()
+
+    torch.manual_seed(23)
+    input_data = np.random.random(input_shape).astype('float32')
+
+    converter_spy = mocker.spy(ModelBuilder, "finish")
+
+    convert_run_compare(edge_program, input_data, tflite_output_preprocess=ToNCHWPreprocess())
+
+    tflite_model = converter_spy.spy_return
+    ops = tflite_model.sub_graphs[0].operators.vector
+    assert len(ops) == 3
+    assert isinstance(ops[0].builtin_options, Reshape)
+    assert isinstance(ops[1].builtin_options, Transpose)
+    assert isinstance(ops[2].builtin_options, Conv2D)
+
+
+def test__formatless_to_formatless(mocker):
+    input_shape = [12, 32]
+    new_shape = [2, 4, 6, 8]
+
+    torch_model = FormatlessToFormatlessModule(new_shape=new_shape)
+    edge_program = to_edge_program(torch_model, input_shape).exported_program()
+
+    torch.manual_seed(23)
+    input_data = np.random.random(input_shape).astype('float32')
+
+    converter_spy = mocker.spy(ModelBuilder, "finish")
+
+    convert_run_compare(edge_program, input_data)
+
+    tflite_model = converter_spy.spy_return
+    ops = tflite_model.sub_graphs[0].operators.vector
+    assert len(ops) == 1  # No extra Transpose ops.
+    assert isinstance(ops[0].builtin_options, Reshape)