COMPMID-3324: Allow in-place mutation on nodes with separate multiple inputs

The current implementation of the in-place mutator is limited by the
assumption that each node only produces one output and therefore all
its output edges represent the same tensor. This means that doing
in-place computations on it, would modify the input of other nodes.

However, this is not the case for operators like split, which creates N
different tensors, one for each output edge.

This patches relaxes this limitation by checking whether the tensor
related to a specific output edge is the same as any of the other edges.
If this is the case, then in-place computation cannot be done.

Change-Id: I3302b41b1f6ec3b4f2ac425ba381778f1c0a4f31
Signed-off-by: Michele Di Giorgio <michele.digiorgio@arm.com>
Reviewed-on: https://review.mlplatform.org/c/ml/ComputeLibrary/+/3421
Comments-Addressed: Arm Jenkins <bsgcomp@arm.com>
Tested-by: Arm Jenkins <bsgcomp@arm.com>
Reviewed-by: Manuel Bottini <manuel.bottini@arm.com>
Reviewed-by: Georgios Pinitas <georgios.pinitas@arm.com>

1 files changed, 42 insertions, 1 deletions

diff --git a/src/graph/mutators/InPlaceOperationMutator.cpp b/src/graph/mutators/InPlaceOperationMutator.cpp
index 394dba84ff..7c75149eb6 100644
--- a/src/graph/mutators/InPlaceOperationMutator.cpp
+++ b/src/graph/mutators/InPlaceOperationMutator.cpp
@@ -30,6 +30,47 @@ namespace arm_compute
 {
 namespace graph
 {
+namespace
+{
+// Check if the output edges of the parent node are separate tensors. If not,
+// it means the same output is connected to multiple nodes and computations on
+// these nodes cannot be done in-place.
+bool output_edges_are_separate_tensors(Graph &g, const Edge *input_edge)
+{
+    const auto parent_node   = input_edge->producer();
+    const auto input_tensor  = input_edge->tensor();
+    const auto input_edge_id = input_edge->id();
+
+    if(parent_node == nullptr)
+    {
+        return false;
+    }
+
+    const auto output_edges = parent_node->output_edges();
+
+    // If the output is connected to only one edge, then computations can
+    // be done in-place.
+    if(output_edges.size() == 1)
+    {
+        return true;
+    }
+
+    return std::all_of(output_edges.begin(),
+                       output_edges.end(),
+                       [&](const EdgeID & edge_id)
+    {
+        // Skip check on current input edge
+        if(edge_id == input_edge_id)
+        {
+            return true;
+        }
+
+        auto edge = g.edge(edge_id);
+        return edge->tensor() != input_tensor;
+    });
+}
+} // namespace
+
 const char *InPlaceOperationMutator::name()
 {
     return "InPlaceOperationMutator";
@@ -60,7 +101,7 @@ void InPlaceOperationMutator::mutate(Graph &g)
             Edge *input_edge = node->input_edge(0);
 
             // Check if parent has a single output if yes then force in place calculation else not
-            if((input_edge != nullptr) && (input_edge->producer() != nullptr) && (input_edge->producer()->output_edges().size() == 1))
+            if((input_edge != nullptr) && output_edges_are_separate_tensors(g, input_edge))
             {
                 // Get current and new output tensors
                 auto current_output_tensor = node->output(0);