diff options
author | Jacob Bohlin <jacob.bohlin@arm.com> | 2020-08-13 13:24:59 +0200 |
---|---|---|
committer | Jacob Bohlin <jacob.bohlin@arm.com> | 2020-11-04 08:57:06 +0100 |
commit | 268394d797db60d07eeace05a2c57e927da0ea15 (patch) | |
tree | 73e23e7b86b442367fcb4b529704e7889779ab2f /ethosu/vela/compiler_driver.py | |
parent | 34359589655e62d0f2cadde53320de71add04e94 (diff) | |
download | ethos-u-vela-268394d797db60d07eeace05a2c57e927da0ea15.tar.gz |
MLBEDSW-1974: Set Scratch buffers size
Set the actual size of the Scratch and Fast Scratch buffer and remove both
Scratch buffers from the subgraph inputs.
Signed-off-by: Jacob Bohlin <jacob.bohlin@arm.com>
Change-Id: I9e4213f48289d9136cdd4cd43c668d37c6af8530
Diffstat (limited to 'ethosu/vela/compiler_driver.py')
-rw-r--r-- | ethosu/vela/compiler_driver.py | 6 |
1 files changed, 6 insertions, 0 deletions
diff --git a/ethosu/vela/compiler_driver.py b/ethosu/vela/compiler_driver.py index 05bf65a4..1d7521b1 100644 --- a/ethosu/vela/compiler_driver.py +++ b/ethosu/vela/compiler_driver.py @@ -291,6 +291,12 @@ def compiler_driver(nng, arch, options, scheduler_options): npu_serialisation.rewrite_npu_call_ops(nng, root_sg, arch) + # Set Scratch and Fast_scratch Tensor size + if scratch_tens is not None: + scratch_tens.set_all_shapes([root_sg.memory_used_per_type.get(MemType.Scratch, 0)]) + if scratch_fast_tens is not None: + scratch_fast_tens.set_all_shapes([root_sg.memory_used_per_type.get(MemType.Scratch_fast, 0)]) + # Allocate all Cpu constant tensors, this is done last because the Npu-ops # have to be serialized into flash and scratch tensors first tensor_allocation.allocate_tensors( |