44# LICENSE file in the root directory of this source tree.
55
66import numpy as np
7- from executorch .backends .nxp .backend .data_format import NXP_NODE_FORMAT
87
98from executorch .backends .nxp .backend .edge_helper import (
10- get_non_qdq_users ,
119 input_tensor ,
1210 output_tensor ,
1311 tensor_rank ,
1412)
1513from executorch .backends .nxp .backend .ir .converter import quantization_utils
1614from executorch .backends .nxp .backend .ir .converter .conversion .common import OpsList
17- from executorch .backends .nxp .backend .ir .converter .conversion .translator import (
18- apply_permutation_to ,
19- create_channels_first_to_channels_last_permutation ,
20- create_channels_last_to_channels_first_permutation ,
21- )
2215from executorch .backends .nxp .backend .ir .converter .node_converter import (
2316 CustomDelegationOptions ,
2417 NodeConverter ,
2922from executorch .backends .nxp .backend .ir .tflite_generator .builtin_options import (
3023 reshape_options ,
3124)
32- from executorch .backends .nxp .backend .neutron_operator_support import (
33- transposition_is_supported_on_neutron ,
34- )
35- from executorch .backends .nxp .backend .neutron_target_spec import NeutronTargetSpec
36- from executorch .exir .dialects ._ops import ops as exir_ops
3725from torch .fx import Node
38- from torch .fx .passes .infra .partitioner import Partition
3926from torch .nn import Parameter
4027
4128
@@ -58,92 +45,6 @@ def _is_supported_in_IR(
5845
5946 return True
6047
61- @classmethod
62- def supports_partitioning_result (
63- cls ,
64- node : Node ,
65- partition_list : list [Partition ],
66- custom_delegation_options : CustomDelegationOptions ,
67- neutron_target_spec : NeutronTargetSpec ,
68- parameters_mapping : dict [str , Parameter ],
69- ):
70- view_copy_partitions = [
71- partition for partition in partition_list if node in partition .nodes
72- ]
73- assert len (view_copy_partitions ) == 1
74-
75- input_format = node .args [0 ].meta [NXP_NODE_FORMAT ]
76- output_format = node .meta [NXP_NODE_FORMAT ]
77- input_shape = list (node .args [0 ].meta ["val" ].shape )
78- output_shape = list (node .meta ["val" ].shape )
79- to_nchw_perm = create_channels_last_to_channels_first_permutation (
80- len (input_shape ), True
81- )
82- to_nhwc_perm = create_channels_first_to_channels_last_permutation (
83- len (output_shape ), True
84- )
85- channels_last_input_shape = apply_permutation_to (
86- input_shape ,
87- create_channels_first_to_channels_last_permutation (len (input_shape ), True ),
88- )
89-
90- if input_format .is_channels_first () and (not output_format .is_channels_first ()):
91- # The `view_copy` removes node format. Conversion will require the addition of a `Transpose` operator.
92- # Make sure the `Transpose` will be supported.
93-
94- if not transposition_is_supported_on_neutron (
95- channels_last_input_shape , to_nchw_perm , neutron_target_spec
96- ):
97- # The `Transpose` would have to be removed by the `PermuteFullyConnectedWeightsAfterReshape` pass.
98- # Make sure it will be applied.
99- users = get_non_qdq_users (node )
100- if len (users ) != 1 or (linear_node := users [0 ]).target not in [
101- exir_ops .edge .aten .addmm .default ,
102- exir_ops .edge .aten .mm .default ,
103- ]:
104- return False
105-
106- if linear_node not in view_copy_partitions [0 ].nodes :
107- # The `mm` / `addmm` node will not be delegated within this partition.
108- return False
109-
110- # Make sure the specific requirements of the `PermuteFullyConnectedWeightsAfterReshape` are satisfied.
111- weights_index = (
112- 2 if linear_node .target == exir_ops .edge .aten .addmm .default else 1
113- )
114- if not (
115- input_shape [0 ] == output_shape [0 ] # Preserve batch.
116- and len (output_shape ) == 2
117- and output_shape [1 ]
118- == linear_node .args [weights_index ].meta ["val" ].shape [0 ]
119- ):
120- return False
121-
122- elif (
123- not input_format .is_channels_first ()
124- ) and output_format .is_channels_first ():
125- # The `view_copy` introduces node format. Conversion will require the addition of a `Transpose` operator.
126- # Make sure the `Transpose` will be supported.
127- if not transposition_is_supported_on_neutron (
128- output_shape , to_nhwc_perm , neutron_target_spec
129- ):
130- return False
131-
132- elif input_format .is_channels_first () and output_format .is_channels_first ():
133- # The `view_copy` works with the channels first format, so both tensors will end up being transposed.
134- # Make sure these transpositions are supported.
135- if not (
136- transposition_is_supported_on_neutron (
137- channels_last_input_shape , to_nchw_perm , neutron_target_spec
138- )
139- and transposition_is_supported_on_neutron (
140- output_shape , to_nhwc_perm , neutron_target_spec
141- )
142- ):
143- return False
144-
145- return True
146-
14748 @staticmethod
14849 def _safe_compute_flat_size (shape : list [int | str ]) -> int :
14950 """Compute the flat size of a tensor with given shape. Strings and negative dimensions are treated as '1'.
0 commit comments