blob: d3686160c55dd657645d274bbb7c208e7a08e2c4 [file] [log] [blame]
Patrik Gustavsson8f1f9aa2021-06-28 07:41:58 +02001# Copyright (C) 2021 Arm Limited or its affiliates. All rights reserved.
2#
3# SPDX-License-Identifier: Apache-2.0
4#
5# Licensed under the Apache License, Version 2.0 (the License); you may
6# not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8#
9# www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an AS IS BASIS, WITHOUT
13# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14# See the License for the specific language governing permissions and
15# limitations under the License.
16# Description:
17# The TosaSupportedOperators class which is a collection of all supported operators and parameter checks.
18from collections import defaultdict
19
20from .data_type import DataType
21from .operation import Op
22from .supported_operators_util import docstring_format_args
23from .supported_operators_util import list_formatter
24from .tosa_mapping import optype_to_tosa_op_type
25
26
27class TosaSupportedOperators:
28 # TODO currently sparsely populated
29 # Categorised lists of supported operators
30 convolution_ops = set((Op.Conv2DBias,))
Patrik Gustavssondf995102021-08-23 15:33:59 +020031 depthwise_convolution_ops = set((Op.DepthwiseConv2DBias,))
32 convolution_like_ops = convolution_ops | depthwise_convolution_ops
33
34 # TODO depending on what will be committed
Patrik Gustavssonc74682c2021-08-17 14:26:38 +020035 max_pooling_ops = Op.op_set(Op.is_maxpool_op)
36 avg_pooling_ops = Op.op_set(Op.is_avgpool_op)
Patrik Gustavssondf995102021-08-23 15:33:59 +020037 pooling_ops = max_pooling_ops | avg_pooling_ops
38 fc_vector_products = set((Op.FullyConnected,))
Patrik Gustavssonc74682c2021-08-17 14:26:38 +020039
Patrik Gustavssondf995102021-08-23 15:33:59 +020040 mac_main_ops = convolution_like_ops | pooling_ops | fc_vector_products
Patrik Gustavssonf1580f02021-09-01 12:43:02 +020041 memory_only_ops = set((Op.Reshape, Op.Transpose, Op.Concat, Op.SplitSliceRead,))
Patrik Gustavsson8f1f9aa2021-06-28 07:41:58 +020042
Patrik Gustavssonb081d672021-08-25 13:49:25 +020043 binary_elem_wise_add_mul_sub = set((Op.Add, Op.Mul, Op.RescaleMul, Op.Sub,))
44
Patrik Gustavsson8f1f9aa2021-06-28 07:41:58 +020045 type_conversion_ops = set((Op.Rescale,))
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +020046 relu_ops = set((Op.Clamp, Op.ReluN,))
Patrik Gustavsson8f1f9aa2021-06-28 07:41:58 +020047 activation_ops = relu_ops
48
49 npu_post_ops = activation_ops
Patrik Gustavssonb081d672021-08-25 13:49:25 +020050 supported_operators = (
51 mac_main_ops | type_conversion_ops | npu_post_ops | memory_only_ops | binary_elem_wise_add_mul_sub
52 )
Patrik Gustavsson8f1f9aa2021-06-28 07:41:58 +020053
54 # Supported data types
55 # TODO will differ compared to TensorFlow Lite, currently set to the same
Patrik Gustavssonf1580f02021-09-01 12:43:02 +020056 supported_op_dtypes = set((DataType.uint8, DataType.int8, DataType.int16, DataType.int32)) # TODO add bool
Patrik Gustavssonf366fb12021-09-07 13:30:29 +020057 tens_dim_range = (1, 65535) # TODO HW limitation, that is to be resolved in SW
Patrik Gustavsson8f1f9aa2021-06-28 07:41:58 +020058
59 def __init__(self):
60 # Setup the generic constraints. Note: the order matters
61 self.generic_constraints = []
62 self.generic_constraints.append(TosaSupportedOperators.constraint_tens_dtype)
Patrik Gustavssonf366fb12021-09-07 13:30:29 +020063 self.generic_constraints.append(TosaSupportedOperators.constraint_tens_dimension)
Patrik Gustavsson8f1f9aa2021-06-28 07:41:58 +020064
65 # Setup specific constraints. Note: the order matters
66 self.specific_constraints = defaultdict(list)
67
Patrik Gustavssondf995102021-08-23 15:33:59 +020068 self.specific_constraints[Op.Transpose].append(TosaSupportedOperators.constraint_ifm_producer)
69
70 # Depthwise Conv specific checks:
71 for op_type in TosaSupportedOperators.depthwise_convolution_ops:
72 self.specific_constraints[op_type].append(TosaSupportedOperators.constraint_depth_multiplier)
73
Patrik Gustavssonf366fb12021-09-07 13:30:29 +020074 # Avgpool specific checks
75 for op_type in TosaSupportedOperators.avg_pooling_ops:
76 self.specific_constraints[op_type].append(TosaSupportedOperators.constraint_padding)
77
Patrik Gustavsson8f1f9aa2021-06-28 07:41:58 +020078 def is_operator_supported(self, op):
79 ext_type = optype_to_tosa_op_type(op.type)
80 if op.type not in TosaSupportedOperators.supported_operators:
81 if op.type not in (Op.Placeholder, Op.SubgraphInput, Op.Const):
82 print(f"Info: {ext_type} '{op.name}' is not a NPU op")
83 return False
84
85 for constraint in self.generic_constraints + self.specific_constraints[op.type]:
86 valid, extra = constraint(op)
87 if not valid:
88 print(f"Warning: {ext_type} '{op.name}' is not supported on the NPU")
89 print(f" - {constraint.__doc__}")
90 if extra:
91 print(f" {extra}")
92 return False
93
94 return True
95
96 # TODO this function is the same for TensorFlow Lite, but input might differ
97 @classmethod
98 @docstring_format_args([list_formatter(supported_op_dtypes)])
99 def constraint_tens_dtype(cls, op):
100 "Tensors must be of type: {}"
101 valid = True
102 extra = []
103 tensors = [tens for tens in op.get_ifm_ifm2_weights_ofm() if tens]
104 if not tensors:
105 tensors = [tens for tens in op.inputs if tens]
106 for tens in tensors:
107 if tens.dtype not in cls.supported_op_dtypes:
108 valid = False
109 extra.append(f"Tensor '{tens.name}' has data type: {tens.dtype}")
110 return valid, ", ".join(extra)
Patrik Gustavssondf995102021-08-23 15:33:59 +0200111
Patrik Gustavssonf366fb12021-09-07 13:30:29 +0200112 # TODO Duplicates check present for TFLite. But it is only temporarily added
113 # This is for a HW limitation, that is to be resolved in SW later on
114 @classmethod
115 @docstring_format_args(tens_dim_range)
116 def constraint_tens_dimension(cls, op):
117 "Tensor dimensions must be in the range [{}, {}]"
118 tens_min, tens_max = cls.tens_dim_range
119 valid = True
120 extra = []
121 tensors = [tens for tens in op.get_ifm_ifm2_weights_ofm() if tens]
122 if not tensors:
123 tensors = [tens for tens in op.inputs if tens]
124 for tens in tensors:
125 if not all(tens_min <= dim <= tens_max for dim in tens.shape):
126 valid = False
127 extra.append(f"Tensor '{tens.name}' has shape: {tens.shape}")
128 return valid, ", ".join(extra)
129
Patrik Gustavssondf995102021-08-23 15:33:59 +0200130 @staticmethod
131 def constraint_ifm_producer(cls, op):
132 "Input must be constant data"
133 valid = op.ifm.ops and op.ifm.ops[0].type == Op.Const
134 return valid, "Op has ifm with non-constant data"
135
Patrik Gustavssonf366fb12021-09-07 13:30:29 +0200136 @staticmethod
137 def constraint_padding(op):
138 # TODO Only support for when global scaling can be used.
139 # That is when there is padding no padding
140 "Avgpool only supported for no padding"
141 top, left, _, _ = op.attrs["explicit_padding"]
142 valid = top == 0 and left == 0
143
144 return valid, "Avgpool with pad_top {top} and pad_left {left}"
145
146 # TODO duplicates tflite_supported operators, but support for depth multiplier should be added at a later stage
Patrik Gustavssondf995102021-08-23 15:33:59 +0200147 @staticmethod
148 def constraint_depth_multiplier(op):
149 "For depth multipliers > 1, IFM channels must be 1 and OFM channels must be equal to the depth multiplier"
150 depth_multiplier = op.attrs.get("depth_multiplier", 1)
151 if depth_multiplier > 1:
152 ifm_channels = op.ifm.shape[3]
153 ofm_channels = op.ofm.shape[3]
154 valid = (ifm_channels == 1) and (ofm_channels == depth_multiplier)
155 extra = (
156 f"Op has ifm_channels={ifm_channels}, ofm_channels={ofm_channels}"
157 f" and depth_multiplier={depth_multiplier}"
158 )
159 return valid, extra
160 return True, "Op has depth_multiplier=1"