MLBEDSW-6313 Static optimisation for Shape OP

*Shape OP value is available at compile time hence
it can be optimised
*Disconnected shape OP at compile time from parent
tensor
*Transformed shape OP tensor into constant

Change-Id: I0a024269e2b592c6146dd72e62d7a41951fb727a
Signed-off-by: Ayaan Masood <Ayaan.Masood@arm.com>

6 files changed, 90 insertions, 18 deletions

diff --git a/SUPPORTED_OPS.md b/SUPPORTED_OPS.md
index 3e76670f..c258dfbd 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.4.0`
+Vela version: `3.4.0rc3.dev1+g5e0ae55`
 
 This file complies with
 [**Gitiles Markdown syntax**](https://github.com/google/gitiles/blob/master/Documentation/markdown.md)
@@ -55,17 +55,18 @@ Please check the supported operator list for your chosen runtime for further inf
 
 ### TFLite Generic Constraints
 
-This is a list of constraints that all NPU operators must satisfy in order to be scheduled on the NPU.
-
-- Input(s) and Output tensors must not be dynamic
-- Input(s) and Output tensors must have a defined shape
-- Output tensors cannot be scalar
-- Scalar Input tensors are only valid for op type: ADD, EXPAND_DIMS, MAXIMUM, MEAN, MINIMUM, MUL, SPLIT, SPLIT_V, SUB
-- Input(s) and Output tensors must not be greater than 4D
-- Input(s), Output and Weight tensors must have quantization parameters
-- Input(s), Output and Weight tensors with quantization scales must be finite
-- Input and Output tensors must have quantization scales that fit within float32 precision
-- Constant tensors should not have NoneType-values
+This is a list of constraints most NPU operators must satisfy in order to be scheduled on the NPU.
+(Operators excluded from certain constraints are shown in brackets [ ] )
+
+- Input(s) and Output tensors must not be dynamic 
+- Input(s) and Output tensors must have a defined shape 
+- Output tensors cannot be scalar 
+- Scalar Input tensors are only valid for op type: ADD, EXPAND_DIMS, MAXIMUM, MEAN, MINIMUM, MUL, SPLIT, SPLIT_V, SUB 
+- Input(s) and Output tensors must not be greater than 4D 
+- Input(s), Output and Weight tensors must have quantization parameters - [Shape]
+- Input(s), Output and Weight tensors with quantization scales must be finite - [Shape]
+- Input and Output tensors must have quantization scales that fit within float32 precision - [Shape]
+- Constant tensors should not have NoneType-values 
 - Tensors must be of type: int16, int32, int8, uint8
 - Tensors which are int32 are only valid when op type is: ADD, MUL, SUB
 - Tensor dimensions must be in the range [1, 65535]
diff --git a/ethosu/vela/operation.py b/ethosu/vela/operation.py
index 8e06b1ef..6b6671be 100644
--- a/ethosu/vela/operation.py
+++ b/ethosu/vela/operation.py
@@ -261,7 +261,7 @@ class Op(Enum):
     SegmentSum = OperatorInfo()
     Select = OperatorInfo()
     SelectV2 = OperatorInfo()
-    Shape = OperatorInfo()
+    Shape = OperatorInfo(indices=NNG_IFM_INDICES)
     Sigmoid = OperatorInfo(indices=NNG_IFM_INDICES)
     SignBit = OperatorInfo()
     Sin = OperatorInfo()
diff --git a/ethosu/vela/tflite_graph_optimiser.py b/ethosu/vela/tflite_graph_optimiser.py
index 06395784..cf3985e4 100644
--- a/ethosu/vela/tflite_graph_optimiser.py
+++ b/ethosu/vela/tflite_graph_optimiser.py
@@ -1391,12 +1391,43 @@ def convert_mean_to_depthwise_conv_or_avgpool(op, arch, nng):
     return op
 
 
+def convert_shape_op_to_constant_tensor(op: Operation, arch, nng):
+    """Static optimisation for SHAPE operator output value known at compile time"""
+
+    # Disconnect SHAPE operator from its parent and transform SHAPE OP into constant
+
+    if op.type == Op.Shape and op.run_on_npu:
+
+        ifm, ofm = op.get_ifm_ofm()
+
+        if len(ifm.shape) != ofm.shape[0]:
+            return op
+
+        # Remove reference of the current shape op from the parent tensor's consumer list
+        ifm.consumer_list = [consumer for consumer in ifm.consumer_list if consumer.op_index != op.op_index]
+
+        # Clear any references to parent node
+        op.inputs = []
+
+        # Convert this SHAPE op to const
+        op.type = Op.Const
+
+        # Add size calculation to shape output tensors
+        ofm.values = np.array(ifm.shape)
+
+    return op
+
+
 def supported_operator_check(op, arch, nng):
     op.run_on_npu = arch.tflite_supported_operators.is_operator_supported(op)
     return op
 
 
 def tflite_optimise_graph(nng, arch):
+
+    # Compile time optimisations
+    optimisation_list = [convert_shape_op_to_constant_tensor]
+
     # Pre-processing step
     pre_process_list = [
         supported_operator_check,
@@ -1409,6 +1440,16 @@ def tflite_optimise_graph(nng, arch):
             sg,
             arch,
             [],
+            optimisation_list,
+            rewrite_unsupported=False,
+        )
+
+    for idx, sg in enumerate(nng.subgraphs):
+        nng.subgraphs[idx] = rewrite_graph.rewrite_graph_pre_order(
+            nng,
+            sg,
+            arch,
+            [],
             pre_process_list,
             rewrite_unsupported=False,
         )
diff --git a/ethosu/vela/tflite_mapping.py b/ethosu/vela/tflite_mapping.py
index 7b487ae4..bf155b9c 100644
--- a/ethosu/vela/tflite_mapping.py
+++ b/ethosu/vela/tflite_mapping.py
@@ -767,7 +767,7 @@ builtin_operator_map = {
     BuiltinOperator.SHAPE: (
         Op.Shape,
         OptionsSerializer("ShapeOptions", (("out_type", datatype_deserialize, datatype_serialize),)),
-        TFLITE_NO_INDICES,
+        TFLITE_IFM_INDICES,
     ),
     BuiltinOperator.POW: (Op.Pow, OptionsSerializer("PowOptions"), TFLITE_NO_INDICES),
     BuiltinOperator.ARG_MIN: (
diff --git a/ethosu/vela/tflite_model_semantic.py b/ethosu/vela/tflite_model_semantic.py
index e0541df5..ee66d4cc 100644
--- a/ethosu/vela/tflite_model_semantic.py
+++ b/ethosu/vela/tflite_model_semantic.py
@@ -186,7 +186,15 @@ class TFLiteSemantic:
         if op.type in (Op.Placeholder, Op.SubgraphInput, Op.Const):
             return True
 
-        for constraint in self.generic_constraints + self.specific_constraints[op.type]:
+        # Generic constraints list filtered out to exclude certain constraints depending on op.type
+        filtered_generic_constraints = []
+
+        for constraint in self.generic_constraints:
+            # Check constraint not in dictionary otherwise return empty array
+            if constraint not in self.get_generic_constraint_exclude_list().get(op.type, []):
+                filtered_generic_constraints.append(constraint)
+
+        for constraint in filtered_generic_constraints + self.specific_constraints[op.type]:
             valid, extra = constraint(op)
             if not valid:
                 print(
@@ -200,6 +208,19 @@ class TFLiteSemantic:
         return True
 
     @staticmethod
+    def get_generic_constraint_exclude_list():
+
+        # Not all generic constraints can be applied to each operator
+        generic_constraints_exclude_list = {
+            Op.Shape: [
+                TFLiteSemantic.constraint_tens_quant_none_check,
+                TFLiteSemantic.constraint_tens_quant_scale,
+                TFLiteSemantic.constraint_quant_scale_inf,
+            ]
+        }
+        return generic_constraints_exclude_list
+
+    @staticmethod
     def constraint_none_const_tensors(op):
         "Constant tensors should not have NoneType-values"
         valid = True
diff --git a/ethosu/vela/vela.py b/ethosu/vela/vela.py
index 90cff034..1de437bb 100644
--- a/ethosu/vela/vela.py
+++ b/ethosu/vela/vela.py
@@ -250,13 +250,22 @@ def generate_supported_ops():
             "",
             f"### {network_type.name} Generic Constraints",
             "",
-            "This is a list of constraints that all NPU operators must satisfy in order to be scheduled on the NPU.",
-            "",
+            "This is a list of constraints most NPU operators must satisfy in order to be scheduled on the NPU.",
+            "(Operators excluded from certain constraints are shown in brackets [ ] )\n" "",
         ]
         for constraint in semantic_checker.generic_constraints:
             # Markdown needs two spaces at the end of a line to render it as a separate line
             reason = constraint.__doc__.replace("\n", "  \n")
-            lines.append(f"- {reason}")
+
+            exclude_list = TFLiteSemantic.get_generic_constraint_exclude_list().items()
+            constraints_excluded_names = [
+                op.name for op, exclude_constraint in exclude_list if constraint in exclude_constraint
+            ]
+            excluded_constraint_text = ""
+            if constraints_excluded_names:
+                excluded_constraint_text = f"- [{', '.join(constraints_excluded_names)}]"
+
+            lines.append(f"- {reason} {excluded_constraint_text}")
         for constraint in supported.generic_constraints:
             # Markdown needs two spaces at the end of a line to render it as a separate line
             reason = constraint.__doc__.replace("\n", "  \n")