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>
diff --git a/ethosu/vela/scheduler.py b/ethosu/vela/scheduler.py
index c35c156..d51b5ac 100644
--- a/ethosu/vela/scheduler.py
+++ b/ethosu/vela/scheduler.py
@@ -926,6 +926,26 @@
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):