Add in-place computation for elementwise operations

- Add in-place computation for elementwise operations at graph level
- Modify support case to test in-place computation for elementwise operations

Resolves: COMPMID-4414

Signed-off-by: Sheri Zhang <sheri.zhang@arm.com>
Change-Id: I5a4de1235dd29a31160e770a16d62f4b98c84ae6
Reviewed-on: https://review.mlplatform.org/c/ml/ComputeLibrary/+/5803
Comments-Addressed: Arm Jenkins <bsgcomp@arm.com>
Reviewed-by: Michele Di Giorgio <michele.digiorgio@arm.com>
Reviewed-by: SiCong Li <sicong.li@arm.com>
Tested-by: Arm Jenkins <bsgcomp@arm.com>

1 files changed, 77 insertions, 16 deletions

diff --git a/src/graph/mutators/InPlaceOperationMutator.cpp b/src/graph/mutators/InPlaceOperationMutator.cpp
index 61639a8f6f..616ec5c73d 100644
--- a/src/graph/mutators/InPlaceOperationMutator.cpp
+++ b/src/graph/mutators/InPlaceOperationMutator.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2018-2020 Arm Limited.
+ * Copyright (c) 2018-2021 Arm Limited.
  *
  * SPDX-License-Identifier: MIT
  *
@@ -23,6 +23,7 @@
  */
 #include "arm_compute/graph/mutators/InPlaceOperationMutator.h"
 
+#include "arm_compute/core/Validate.h"
 #include "arm_compute/graph/Graph.h"
 #include "arm_compute/graph/Logger.h"
 
@@ -69,6 +70,64 @@ bool output_edges_are_separate_tensors(Graph &g, const Edge *input_edge)
         return edge->tensor() != input_tensor;
     });
 }
+
+// If do in-place calculation, then need to use the new output and inherit original output's accessor
+void set_new_output_and_inherit_accessor(std::unique_ptr<INode> &node, Tensor *orig_output, Tensor *new_output)
+{
+    ARM_COMPUTE_LOG_GRAPH_INFO("Switching to in-place computation for the node with ID : "
+                               << node->id() << " and name : " << node->name() << std::endl);
+    // Update accessor
+    new_output->set_accessor(orig_output->extract_accessor());
+    // Update output
+    node->set_output_tensor(new_output->id(), 0);
+}
+
+// Try to mutate the node to perform the elementwise in-place calculation
+void try_in_place_elementwise(std::unique_ptr<INode> &node)
+{
+    // Get input edge
+    Edge *input0_edge = node->input_edge(0);
+    Edge *input1_edge = node->input_edge(1);
+    ARM_COMPUTE_ERROR_ON(input0_edge == nullptr || input1_edge == nullptr);
+
+    auto input0_tensor = input0_edge->tensor();
+    auto input1_tensor = input1_edge->tensor();
+    ARM_COMPUTE_ERROR_ON(input0_tensor == nullptr || input1_tensor == nullptr);
+
+    const auto shape0 = input0_tensor->desc().shape;
+    const auto shape1 = input1_tensor->desc().shape;
+    const auto qinfo0 = input0_tensor->desc().quant_info;
+    const auto qinfo1 = input1_tensor->desc().quant_info;
+
+    const TensorShape out_shape = TensorShape::broadcast_shape(shape0, shape1);
+    // Inputs are not broadcast compatible
+    if(out_shape.total_size() == 0)
+    {
+        return;
+    }
+
+    // Get current output tensor
+    auto current_output_tensor = node->output(0);
+    ARM_COMPUTE_ERROR_ON(current_output_tensor == nullptr);
+    const auto qinfo_out = current_output_tensor->desc().quant_info;
+
+    // Can do in place, if the input has same shape as output, has same quntisation info as output, and input doesn't have accessor.
+    bool input0_can_in_place = !arm_compute::detail::have_different_dimensions(out_shape, shape0, 0) && (qinfo0 == qinfo_out) && (input0_tensor->accessor() == nullptr);
+    bool input1_can_in_place = !arm_compute::detail::have_different_dimensions(out_shape, shape1, 0) && (qinfo1 == qinfo_out) && (input1_tensor->accessor() == nullptr);
+
+    if(input0_can_in_place)
+    {
+        set_new_output_and_inherit_accessor(node, current_output_tensor, input0_tensor);
+    }
+    else if(input1_can_in_place)
+    {
+        set_new_output_and_inherit_accessor(node, current_output_tensor, input1_tensor);
+    }
+    else
+    {
+        ARM_COMPUTE_LOG_GRAPH_VERBOSE("Prevented in-place operation as there is an accessor bound to the input tensor or the quantization info are different.\n");
+    }
+}
 } // namespace
 
 const char *InPlaceOperationMutator::name()
@@ -103,25 +162,27 @@ void InPlaceOperationMutator::mutate(Graph &g)
             // Check if parent has a single output if yes then force in place calculation else not
             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);
-                auto new_output_tensor     = input_edge->tensor();
-
-                ARM_COMPUTE_ERROR_ON(current_output_tensor == nullptr || new_output_tensor == nullptr);
-
-                // Prevent in-place operation if there is an accessor bound to the in-place tensor or quantization info are different
-                if(new_output_tensor->accessor() != nullptr || current_output_tensor->desc().quant_info != new_output_tensor->desc().quant_info)
+                if(node->type() == NodeType::EltwiseLayer)
                 {
-                    ARM_COMPUTE_LOG_GRAPH_VERBOSE("Prevented in-place operation as there is an accessor bound to the input tensor or the quantization info are different.\n");
+                    try_in_place_elementwise(node);
                 }
                 else
                 {
-                    ARM_COMPUTE_LOG_GRAPH_VERBOSE("Switching to in-place computation for the node with ID : "
-                                                  << node->id() << " and name : " << node->name() << std::endl);
-                    // Update accessor
-                    new_output_tensor->set_accessor(current_output_tensor->extract_accessor());
-                    // Update output
-                    node->set_output_tensor(new_output_tensor->id(), 0);
+                    // Get current and new output tensors
+                    auto current_output_tensor = node->output(0);
+                    auto new_output_tensor     = input_edge->tensor();
+
+                    ARM_COMPUTE_ERROR_ON(current_output_tensor == nullptr || new_output_tensor == nullptr);
+
+                    // Prevent in-place operation if there is an accessor bound to the in-place tensor or quantization info are different
+                    if(new_output_tensor->accessor() != nullptr || current_output_tensor->desc().quant_info != new_output_tensor->desc().quant_info)
+                    {
+                        ARM_COMPUTE_LOG_GRAPH_VERBOSE("Prevented in-place operation as there is an accessor bound to the input tensor or the quantization info are different.\n");
+                    }
+                    else
+                    {
+                        set_new_output_and_inherit_accessor(node, current_output_tensor, new_output_tensor);
+                    }
                 }
             }
         }