From 107063615e46a9d9d0f07db9cf139ac136978c77 Mon Sep 17 00:00:00 2001 From: Johan Alfven Date: Thu, 13 Apr 2023 12:20:55 +0200 Subject: MLBEDSW-7503: Avoid changing buffer index for models with only CPU ops - When compiling a model that only contains CPU ops, Vela unnecessary adds an empty buffer. - This extra buffer is added due to that the fast scratch tensor always occupies index 1. - Since scratch and fast_scratch does not have any constant data they can use buffer 0. Change-Id: I25e1fb124deed7069641bde1f571b522c5bf763a Signed-off-by: Johan Alfven --- ethosu/vela/tflite_writer.py | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/ethosu/vela/tflite_writer.py b/ethosu/vela/tflite_writer.py index c8250c6e..2e7345ce 100644 --- a/ethosu/vela/tflite_writer.py +++ b/ethosu/vela/tflite_writer.py @@ -75,9 +75,10 @@ def make_vector(v): class TFLiteSerialiser: - BUF_IDX_SCRATCH = 0 # Always assign scratch to buffer 0 - BUF_IDX_SCRATCH_FAST = 1 # Always assign scratch_fast to buffer 1 - BUF_IDX_START = 2 # Unique buffer id for every tensor in all subgraphs + # The 0th buffer is always by default an empty buffer that can be used by tensors + # without any constant data + BUF_IDX_ZERO = 0 + BUF_IDX_START = 1 def __init__(self, nng): self.builder = flatbuffers.Builder(0) @@ -161,11 +162,10 @@ class TFLiteSerialiser: for tens in tensors: # Set buffer ids depending on allocation - if tens.is_allocated_in_tensor_arena(scratch_tensor_mem_area): - buffer_map[tens] = TFLiteSerialiser.BUF_IDX_SCRATCH - elif tens.mem_type == MemType.Scratch_fast: - # For Scratch_fast when not co-allocated with scratch in the TensorArena: - buffer_map[tens] = TFLiteSerialiser.BUF_IDX_SCRATCH_FAST + if tens.is_allocated_in_tensor_arena(scratch_tensor_mem_area) or tens.mem_type == MemType.Scratch_fast: + # Tensor allocated in the scratch areas, does not have any constant data and can + # therefore all point to the empty buffer (zero) + buffer_map[tens] = TFLiteSerialiser.BUF_IDX_ZERO else: buffer_map[tens] = self.buf_idx self.buf_idx += 1 @@ -256,6 +256,8 @@ class TFLiteSerialiser: values = tens.values buf_id = self.buffer_map[tens] + # Sanity check that if buffer 0 is used there must not be any data + assert not (buf_id == TFLiteSerialiser.BUF_IDX_ZERO and values is not None) self.buffers_to_write[buf_id] = None if values is None else values.flatten().view(np.uint8) shape = self.write_int_vector(tens_shape) -- cgit v1.2.1