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 <johan.alfven@arm.com>

1 files 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)