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
diff --git a/ethosu/vela/compiler_driver.py b/ethosu/vela/compiler_driver.py
index 05bf65a..1d7521b 100644
--- a/ethosu/vela/compiler_driver.py
+++ b/ethosu/vela/compiler_driver.py
@@ -291,6 +291,12 @@
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(