MLBEDSW-4976: index errors in scheduler

- Fixed index error in memory_snapshot
- When removing a cascade, also references are removed

Change-Id: I2b35dc52671d8ce115eb32bfdd93584391d1fc6d
Signed-off-by: Louis Verhaard <louis.verhaard@arm.com>

1 files changed, 6 insertions, 4 deletions

diff --git a/ethosu/vela/scheduler.py b/ethosu/vela/scheduler.py
index a375a939..de10bad7 100644
--- a/ethosu/vela/scheduler.py
+++ b/ethosu/vela/scheduler.py
@@ -861,10 +861,12 @@ class Scheduler:
             return max_sched
 
         # Extract the cascades
-        cascades = [cascade for cascade in schedule.cascades.values()]
-        for cascade_info in cascades:
-            # Remove existing cascade from schedule
-            del schedule.cascades[cascade_info.end]
+        cascades = schedule.cascades
+        # Remove existing cascade from schedule
+        schedule.cascades = {}
+        for cost in schedule.cost_map.values():
+            cost.cascade = 0
+        for cascade_info in cascades.values():
             # Optimize the sub-schedule in this cascade
             opt_sub_schedule = self.optimize_sub_schedule(cascade_info, schedule, max_template, sram_limit)
             # Update the sub-schedule Op and cascade costs to the full schedule