MLBEDSW-5384 FC layers run on NPU if underlying shape is 2D
*Added generic function which checks if underlying shape of
FullyConnected operation is 2D and performs shape reduction
*Fully connected operation >2 dimensions now run on NPU if the above
case is satisfied
*constraint_fc_output_2d and rewrite_fully_connected_input refactored
*Added unit test to confirm this functionality
Signed-off-by: Ayaan Masood <Ayaan.Masood@arm.com>
Change-Id: I0e29c767e5b84841eb53bbc44464b36a454f7b38
diff --git a/ethosu/vela/tflite_model_semantic.py b/ethosu/vela/tflite_model_semantic.py
index b264479..c811a0d 100644
--- a/ethosu/vela/tflite_model_semantic.py
+++ b/ethosu/vela/tflite_model_semantic.py
@@ -295,14 +295,11 @@
@staticmethod
def constraint_fc_output_2d(op):
- "The output tensor(s) must have 2D shape"
- valid = True
- extra = []
- for tens in op.outputs:
- if len(tens.shape) != 2:
- valid = False
- extra.append(f"Tensor '{tens.name}' is {len(tens.shape)}D")
- return valid, ", ".join(extra)
+ """The output tensor(s) must have 2D shape"""
+ valid = op.ifm.get_shape_as_2d(op.weights.shape[-2]) is not None
+ extra = f"Op has non-2D output tensor '{op.ofm.name}'" if not valid else ""
+
+ return valid, extra
@staticmethod
def constraint_stride_type(op):