MLBEDSW-6148: Allow overwrite of subgraph input

This change will allow the subgraph's input tensor
to be reused/overwritten by the output from an elementwise op
if there is only one consumer attached to the input tensor.

Signed-off-by: Johan Alfven <johan.alfven@arm.com>
Change-Id: I317188af11a5470614770e18dc8973462fd5f21c
diff --git a/ethosu/vela/extract_npu_subgraphs.py b/ethosu/vela/extract_npu_subgraphs.py
index f46f031..ac24e43 100644
--- a/ethosu/vela/extract_npu_subgraphs.py
+++ b/ethosu/vela/extract_npu_subgraphs.py
@@ -87,6 +87,14 @@
         call_ps.inputs.append(orig_tens)
         call_ps.primary_op.inputs.append(orig_tens)
 
+    # Elementwise op can not overwrite ifm if input is used by many consumers
+    if orig_tens in cpu_subgraph.input_tensors and len(orig_tens.consumers()) > 1:
+        new_tens.ifm_write_protected = True
+
+    # Elementwise op can not overwrite ifm if tensor is used as output from sub graph
+    if orig_tens in cpu_subgraph.output_tensors:
+        new_tens.ifm_write_protected = True
+
     for op in list(orig_tens.consumers()):
         if op is None:
             continue  # Subgraph consumers handled separately.
diff --git a/ethosu/vela/live_range.py b/ethosu/vela/live_range.py
index 5aec0df..1aaaadd 100644
--- a/ethosu/vela/live_range.py
+++ b/ethosu/vela/live_range.py
@@ -158,7 +158,7 @@
 
 def merge_elementwise_op_ranges(sg, sched_op, lr_graph, target_mem_area, target_mem_type_set):
     def _tensor_should_be_ignored(tens):
-        if tens in sg.input_tensors + sg.output_tensors:
+        if tens.ifm_write_protected:
             return True
         return tensor_should_be_ignored(tens, target_mem_area, target_mem_type_set)
 
diff --git a/ethosu/vela/tensor.py b/ethosu/vela/tensor.py
index ab63afd..82de897 100644
--- a/ethosu/vela/tensor.py
+++ b/ethosu/vela/tensor.py
@@ -370,6 +370,7 @@
         "resampling_mode",
         "src_tensor",
         "needs_linear_format",
+        "ifm_write_protected",
     )
     AllocationQuantum = 16
 
@@ -416,6 +417,7 @@
         self.resampling_mode: resampling_mode = resampling_mode.NONE
 
         self.needs_linear_format = True
+        self.ifm_write_protected = False
 
         # Reference to parent-tensor if this tensor is a clone
         self.src_tensor = None