COMPMID-920: Introduce prepare() stage

Change-Id: I08ddb7f6e061178e7566518b48e4e18f8f078596
Reviewed-on: https://eu-gerrit-1.euhpc.arm.com/129825
Tested-by: Jenkins <bsgcomp@arm.com>
Reviewed-by: Anthony Barbier <anthony.barbier@arm.com>

3 files changed, 104 insertions, 15 deletions

diff --git a/src/graph/GraphManager.cpp b/src/graph/GraphManager.cpp
index c0720ac685..fa7dfdf8f8 100644
--- a/src/graph/GraphManager.cpp
+++ b/src/graph/GraphManager.cpp
@@ -74,23 +74,47 @@ void GraphManager::finalize_graph(Graph &graph, GraphContext &ctx, PassManager &
     auto workload = detail::configure_all_nodes(graph, ctx);
     ARM_COMPUTE_ERROR_ON_MSG(workload.tasks.empty(), "Could not configure all nodes!");
 
-    // Allocate all tensors
-    detail::allocate_all_tensors(graph);
+    // TODO (COMPMID-920) : Update prepare for NEON/GC
+    if(forced_target == Target::CL)
+    {
+        // Allocate const tensors and call accessors
+        detail::allocate_const_tensors(graph);
+        detail::call_all_const_node_accessors(graph);
+
+        // Prepare graph
+        detail::prepare_all_tasks(workload);
+
+        // Allocate all tensors
+        detail::allocate_all_tensors(graph);
+
+        // Finalize Graph context
+        ctx.finalize();
 
-    // Call accessors on all Const nodes
-    detail::call_all_const_node_accessors(graph);
+        // Register graph
+        _workloads.insert(std::make_pair(graph.id(), std::move(workload)));
+        ARM_COMPUTE_LOG_GRAPH_VERBOSE("Created workload for graph with ID : " << graph.id().get() << std::endl);
+    }
+    else
+    {
+        // Allocate all tensors
+        detail::allocate_all_tensors(graph);
 
-    _workloads.insert(std::make_pair(graph.id(), std::move(workload)));
-    ARM_COMPUTE_LOG_GRAPH_VERBOSE("Created workload for graph with ID : " << graph.id().get() << std::endl);
+        // Call accessors on all Const nodes
+        detail::call_all_const_node_accessors(graph);
 
-    // Finalize Graph context
-    ctx.finalize();
+        // Finalize Graph context
+        ctx.finalize();
 
-    // Make first run
-    execute_graph(graph);
+        // Register graph
+        _workloads.insert(std::make_pair(graph.id(), std::move(workload)));
+        ARM_COMPUTE_LOG_GRAPH_VERBOSE("Created workload for graph with ID : " << graph.id().get() << std::endl);
 
-    // Release all unused const nodes
-    detail::release_unused_tensors(graph);
+        // Make first run
+        execute_graph(graph);
+
+        // Release all unused const tensors
+        detail::release_unused_tensors(graph);
+    }
 }
 
 void GraphManager::execute_graph(Graph &graph)
diff --git a/src/graph/Workload.cpp b/src/graph/Workload.cpp
index c53a8a42da..f350bbf625 100644
--- a/src/graph/Workload.cpp
+++ b/src/graph/Workload.cpp
@@ -37,5 +37,13 @@ void ExecutionTask::operator()()
         task->run();
     }
 }
+
+void ExecutionTask::prepare()
+{
+    if(task)
+    {
+        task->prepare();
+    }
+}
 } // namespace graph
 } // namespace arm_compute
\ No newline at end of file
diff --git a/src/graph/detail/ExecutionHelpers.cpp b/src/graph/detail/ExecutionHelpers.cpp
index 5a50728164..0bb47f2b33 100644
--- a/src/graph/detail/ExecutionHelpers.cpp
+++ b/src/graph/detail/ExecutionHelpers.cpp
@@ -61,15 +61,61 @@ void configure_all_tensors(Graph &g)
     }
 }
 
+void allocate_all_input_tensors(INode &node)
+{
+    for(unsigned int i = 0; i < node.num_inputs(); ++i)
+    {
+        Tensor *tensor = node.input(i);
+        if(tensor != nullptr && !tensor->bound_edges().empty())
+        {
+            ARM_COMPUTE_ERROR_ON_MSG(!tensor->handle(), "Tensor handle is not configured!");
+            tensor->handle()->allocate();
+        }
+    }
+}
+
+void allocate_all_output_tensors(INode &node)
+{
+    for(unsigned int i = 0; i < node.num_outputs(); ++i)
+    {
+        Tensor *tensor = node.output(i);
+        if(tensor != nullptr && !tensor->bound_edges().empty())
+        {
+            ARM_COMPUTE_ERROR_ON_MSG(!tensor->handle(), "Tensor handle is not configured!");
+            tensor->handle()->allocate();
+        }
+    }
+}
+
+void allocate_const_tensors(Graph &g)
+{
+    for(auto &node : g.nodes())
+    {
+        if(node != nullptr)
+        {
+            switch(node->type())
+            {
+                case NodeType::Const:
+                case NodeType::Input:
+                    allocate_all_output_tensors(*node);
+                    break;
+                case NodeType::Output:
+                    allocate_all_input_tensors(*node);
+                default:
+                    break;
+            }
+        }
+    }
+}
+
 void allocate_all_tensors(Graph &g)
 {
     auto &tensors = g.tensors();
 
     for(auto &tensor : tensors)
     {
-        if(tensor && !tensor->bound_edges().empty())
+        if(tensor && !tensor->bound_edges().empty() && tensor->handle() != nullptr && tensor->handle()->tensor().info()->is_resizable() && tensor->handle()->tensor().is_used())
         {
-            ARM_COMPUTE_ERROR_ON_MSG(!tensor->handle(), "Tensor handle is not configured!");
             tensor->handle()->allocate();
         }
     }
@@ -96,7 +142,8 @@ void validate_all_nodes(Graph &g)
 ExecutionWorkload configure_all_nodes(Graph &g, GraphContext &ctx)
 {
     ExecutionWorkload workload;
-    auto             &nodes = g.nodes();
+    workload.graph = &g;
+    auto &nodes    = g.nodes();
 
     // Create tasks
     for(auto &node : nodes)
@@ -176,6 +223,16 @@ void call_all_input_node_accessors(ExecutionWorkload &workload)
     }
 }
 
+void prepare_all_tasks(ExecutionWorkload &workload)
+{
+    ARM_COMPUTE_ERROR_ON(workload.graph == nullptr);
+    for(auto &task : workload.tasks)
+    {
+        task.prepare();
+        release_unused_tensors(*workload.graph);
+    }
+}
+
 void call_all_tasks(ExecutionWorkload &workload)
 {
     for(auto &task : workload.tasks)