commit 6f4955aa7097b123bbf31aae4654547bb3e3c68c
author Samuel Panijel <samuel.panijel@arm.com> Thu Jun 10 13:40:03 2021 +0300
committer tim.hall <tim.hall@arm.com> Mon Jul 05 11:01:19 2021 +0000
tree 72972b12a5e53743844bbc2d86c3fe50ae4fa425
parent b279844cae93353410210597535fdbf495cd0498

MLBEDSW-3890 handling scratch tensor

vela: Possible issue with handling scratch tensor on non-ethosu custom op

Fixing a case where a tensor input name ends with "scratch".
4 test cases passing this change:
1) non-optimized tflite - input tensor name is _split_1_scratch
2) optimized tflite - input tensor name is _split_1_scratch
3) optimized tflite - input tensor name is _split_1_scratch and custom
   operation name is non_ethus_u
4) non-optimized tflite - input tensor name is _split_1_scratch_fast

Change-Id: Ia515805825b7f9a646607c5075b7ea3a0cf6aad8
Signed-off-by: Samuel Panijel <samuel.panijel@arm.com>

ethosu/vela/npu_serialisation.py

diff --git a/ethosu/vela/npu_serialisation.py b/ethosu/vela/npu_serialisation.py
index f462168..06ea61d 100644
--- a/ethosu/vela/npu_serialisation.py
+++ b/ethosu/vela/npu_serialisation.py
@@ -82,7 +82,7 @@
         sg.scratch_fast_tensor = make_memory_tensor(
             sg.name + "_scratch_fast", scratch_fast_area, MemType.Scratch_fast, 0, False, arch
         )
-        sg.scratch_fast_tensor.purpose = TensorPurpose.Scratch
+        sg.scratch_fast_tensor.purpose = TensorPurpose.ScratchFast
     else:
         sg.scratch_tensor = scratch_tens
         sg.scratch_tensor.shape[0] += scratch_size

ethosu/vela/tflite_writer.py

diff --git a/ethosu/vela/tflite_writer.py b/ethosu/vela/tflite_writer.py
index 687b887..8cabb0a 100644
--- a/ethosu/vela/tflite_writer.py
+++ b/ethosu/vela/tflite_writer.py
@@ -25,6 +25,7 @@
 from .nn_graph import PassPlacement
 from .operation import Op
 from .tensor import MemType
+from .tensor import TensorPurpose
 from .tflite import Buffer
 from .tflite import Metadata
 from .tflite import Model
@@ -37,6 +38,7 @@
 from .tflite_mapping import BuiltinOperator
 from .tflite_mapping import datatype_inv_map
 
+
 # ugh, the python flatbuffer interface is missing a method to add in file identifier. patching it in here:
 
 tflite_version = 3
@@ -345,7 +347,7 @@
 
         all_tensors = [tens for nm, idx, tens in sorted((tens.name, idx, tens) for idx, tens in enumerate(tensor_set))]
 
-        scratch_tensors = [tens for tens in all_tensors if tens.name.endswith("scratch")]
+        scratch_tensors = [tens for tens in all_tensors if tens.purpose is TensorPurpose.Scratch]
 
         if len(scratch_tensors) == 0:
             scratch_tensor = None