MLBEDSW-7396: MLCE: Add num elements constraint on reshape

Adding constraint for faulty reshape operators. Number of elements
for IFM and OFM must be the same.

Change-Id: I2e31e9d1e39b5aa3a0c595032a66e14374a0719e
Signed-off-by: Johan Alfven <johan.alfven@arm.com>

2 files changed, 13 insertions, 1 deletions

diff --git a/SUPPORTED_OPS.md b/SUPPORTED_OPS.md
index aff1d1d6..860d1fe6 100644
--- a/SUPPORTED_OPS.md
+++ b/SUPPORTED_OPS.md
@@ -1,7 +1,7 @@
 # Supported Ops
 
 This file was automatically generated by Vela using the `--supported-ops-report` parameter.  
-Vela version: `3.7.0`
+Vela version: `3.7.1.dev2+g19f8967.d20230301`
 
 This file complies with
 [**Gitiles Markdown syntax**](https://github.com/google/gitiles/blob/master/Documentation/markdown.md)
@@ -155,6 +155,7 @@ This is a list of constraints that the DEPTHWISE_CONV_2D operator must satisfy i
 This is a list of constraints that the EXPAND_DIMS operator must satisfy in order to be scheduled on the NPU.
 
 - Input and output quantisation must match.
+- Input and output number of elements must match.
 
 ### TFLite FULLY_CONNECTED Constraints
 
@@ -251,6 +252,7 @@ This is a list of constraints that the PAD operator must satisfy in order to be
 This is a list of constraints that the RESHAPE operator must satisfy in order to be scheduled on the NPU.
 
 - Input and output quantisation must match.
+- Input and output number of elements must match.
 - Shape must be constant
 
 ### TFLite RESIZE_BILINEAR Constraints
@@ -304,6 +306,7 @@ This is a list of constraints that the SPLIT_V operator must satisfy in order to
 This is a list of constraints that the SQUEEZE operator must satisfy in order to be scheduled on the NPU.
 
 - Input and output quantisation must match.
+- Input and output number of elements must match.
 
 ### TFLite STRIDED_SLICE Constraints
 
diff --git a/ethosu/vela/tflite_model_semantic.py b/ethosu/vela/tflite_model_semantic.py
index 0c2086c3..9f53a1e6 100644
--- a/ethosu/vela/tflite_model_semantic.py
+++ b/ethosu/vela/tflite_model_semantic.py
@@ -27,6 +27,7 @@ from .operation import Op
 from .supported_operators_util import docstring_format_args
 from .supported_operators_util import list_formatter
 from .tensor import check_quantized_tens_scaling_equal
+from .tensor import shape_num_elements
 from .tflite_mapping import BUILTIN_OPERATOR_UNKNOWN
 from .tflite_mapping import optype_to_builtintype
 
@@ -148,6 +149,7 @@ class TFLiteSemantic:
         # Ops reshaping dimensions: Reshape, Squeeze and ExpandDims
         for op_type in TFLiteSemantic.reshape_ops:
             self.specific_constraints[op_type].append(TFLiteSemantic.constraint_matching_in_out_quant)
+            self.specific_constraints[op_type].append(TFLiteSemantic.constraint_matching_in_out_elements)
 
         # Softmax specific checks:
         self.specific_constraints[Op.Softmax].append(TFLiteSemantic.constraint_matching_shapes)
@@ -664,6 +666,13 @@ class TFLiteSemantic:
             return False, "IFM and OFM quantisation parameters are not equal."
         return True, "IFM and OFM quantisation parameters matches."
 
+    @staticmethod
+    def constraint_matching_in_out_elements(op):
+        "Input and output number of elements must match."
+        if shape_num_elements(op.ifm.shape) != shape_num_elements(op.ofm.shape):
+            return False, f"IFM {op.ifm.shape} and OFM {op.ofm.shape} number of elements are not equal."
+        return True, "IFM and OFM number of elements are equal."
+
 
 def tflite_semantic_checker(nng):
     semantic_checker = TFLiteSemantic()