diff options
author | Patrik Gustavsson <patrik.gustavsson@arm.com> | 2020-06-29 17:36:55 +0200 |
---|---|---|
committer | tim.hall <tim.hall@arm.com> | 2020-07-07 15:23:20 +0000 |
commit | 3ab9452881a15d88710f4b5d7c14ba5069e74948 (patch) | |
tree | 5f942504d52094f728d974c5bb66ae67aa83eedb /ethosu/vela/tflite_writer.py | |
parent | e843d3311b8945baa32654af0dccb229b6861438 (diff) | |
download | ethos-u-vela-3ab9452881a15d88710f4b5d7c14ba5069e74948.tar.gz |
MLBEDSW-2551 Add support for more mem-cfgs
Added support for one more memory configuration-
Change-Id: Iac19992386e3e9b80bd519acb1b0a399c47d736f
Signed-off-by: Patrik Gustavsson <patrik.gustavsson@arm.com>
Diffstat (limited to 'ethosu/vela/tflite_writer.py')
-rw-r--r-- | ethosu/vela/tflite_writer.py | 15 |
1 files changed, 12 insertions, 3 deletions
diff --git a/ethosu/vela/tflite_writer.py b/ethosu/vela/tflite_writer.py index 4aa23b5f..cf40b5b5 100644 --- a/ethosu/vela/tflite_writer.py +++ b/ethosu/vela/tflite_writer.py @@ -142,7 +142,7 @@ class TFLiteSerialiser: buffer_map = {} - buf_idx = 1 + buf_idx = 2 for tens in tensors: # Set buffer ids depending on allocation @@ -314,7 +314,11 @@ class TFLiteSerialiser: 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.purpose == TensorPurpose.Scratch] + scratch_tensors = [tens for tens in all_tensors if tens.name.endswith("scratch")] + + for tens in all_tensors: + if tens.name.endswith("scratch_fast"): + scratch_fast_tensor = tens if len(scratch_tensors) == 0: scratch_tensor = None @@ -331,11 +335,16 @@ class TFLiteSerialiser: assert all(inp in sg.original_inputs for inp in sg.input_tensors) inputs = [self.tensor_map[tens] for tens in sg.original_inputs] - # Add the Scratch Tensor as input to the NPU subgraph to get it allocated by TensorFlow Lite Micro + # Add the Scratch Tensors as input to the NPU subgraph to get them allocated by TensorFlow Lite Micro scratch_tensor_idx = self.tensor_map.get(scratch_tensor, None) + scratch_fast_tensor_idx = self.tensor_map.get(scratch_fast_tensor, None) + if scratch_tensor_idx is not None and scratch_tensor_idx not in inputs: inputs.append(scratch_tensor_idx) + if scratch_fast_tensor_idx is not None and scratch_fast_tensor_idx not in inputs: + inputs.append(scratch_fast_tensor_idx) + inputs_offset = self.write_int_vector(inputs) outputs_offset = self.write_int_vector([self.tensor_map[tens] for tens in sg.output_tensors]) |