MLBEDSW-1370: Use NHCWB16 between NPU ops

Added support for using NHCWB16 between cascaded passes.
(For Reshape format is kept to NHWC)

Change-Id: I0ef1631984fec89fe09999b64ae69563e2aefc9b
Signed-off-by: Patrik Gustavsson <patrik.gustavsson@arm.com>

1 files changed, 20 insertions, 0 deletions

diff --git a/ethosu/vela/scheduler.py b/ethosu/vela/scheduler.py
index c35c1566..d51b5ac6 100644
--- a/ethosu/vela/scheduler.py
+++ b/ethosu/vela/scheduler.py
@@ -926,6 +926,26 @@ class DynamicProgrammingScheduler:
         self.sg.cascaded_passes = cascaded_passes
         self.sg.build_cascaded_pass_links()
 
+        # Check if NHCWB16 can be used in between cascaded passes
+        # (NHCWB16 within cascaded passes has been handled earlier in this function)
+        if self.sg.placement == PassPlacement.Npu:
+            for ps in self.sg.cascaded_passes:
+                if ps.placement != PassPlacement.Npu:
+                    continue
+                for output in ps.outputs:
+                    if output.purpose != TensorPurpose.FeatureMap:
+                        continue
+
+                    use_NHCWB16 = True
+                    for op in output.consumer_list:
+                        if op == None or op.type == 'Reshape':
+                            use_NHCWB16 = False
+                        else:
+                            use_NHCWB16 &= op.run_on_npu
+
+                    if use_NHCWB16:
+                        output.set_format(TensorFormat.NHCWB16, arch)
+
 
 def schedule_passes(nng, arch, options: SchedulerOptions):