diff options
| -rw-r--r-- | ethosu/vela/cascade_builder.py | 6 | ||||
| -rw-r--r-- | ethosu/vela/scheduler.py | 5 |
2 files changed, 9 insertions, 2 deletions
diff --git a/ethosu/vela/cascade_builder.py b/ethosu/vela/cascade_builder.py index 7baa3a83..ba210032 100644 --- a/ethosu/vela/cascade_builder.py +++ b/ethosu/vela/cascade_builder.py @@ -105,7 +105,11 @@ class CascadeBuilder: def _estimate_sram_usage(self, sched_op, cost) -> int: """Estimate the SRAM required for the Op if all FeatureMaps are in SRAM""" - ifm2_size = sched_op.ifm2_size_in_bytes() + if sched_op.parent_op.type.is_binary_elementwise_op(): + # ifm2 is scalar or constant and will always persist in permanent memory + ifm2_size = 0 + else: + ifm2_size = sched_op.ifm2_size_in_bytes() if sched_op.requires_full_ifm: ifm_size = sched_op.ifm_size_in_bytes() else: diff --git a/ethosu/vela/scheduler.py b/ethosu/vela/scheduler.py index 73eb8b42..c7d08fbb 100644 --- a/ethosu/vela/scheduler.py +++ b/ethosu/vela/scheduler.py @@ -57,6 +57,7 @@ from .nn_graph import Pass from .nn_graph import PassPlacement from .nn_graph import SchedulingStrategy from .nn_graph import Subgraph +from .live_range import ofm_can_reuse_ifm from .numeric_util import round_down from .numeric_util import round_up from .operation import NpuBlockType @@ -974,9 +975,11 @@ class Scheduler: op_mem_usage = 0 else: # Min schedule only have ifm and ofm in SRAM (no buffered weigth tensors) - op_mem_usage = sched_op.ifm_size_in_bytes() + sched_op.ofm_size_in_bytes() + ofm_size = 0 if ofm_can_reuse_ifm(sched_op) else sched_op.ofm_size_in_bytes() + op_mem_usage = sched_op.ifm_size_in_bytes() + ofm_size non_local_mem_usage[sched_op] = min_schedule.memory_snapshot[time_index] - op_mem_usage + assert non_local_mem_usage[sched_op] >= 0 # Crate cascades for Min schedule cascade_builder = CascadeBuilder(self.sched_ops, self.arch.is_spilling_enabled(), non_local_mem_usage) |