From f52e9fd4b42b75916a5e5a777d4ee9c58f721d14 Mon Sep 17 00:00:00 2001
From: mathad01 <matthew.haddon@arm.com>
Date: Wed, 5 May 2021 12:51:56 +0100
Subject: IVGCVSW-5418 ExecuteNetwork test for MobileBERT

 * Fix bug in stack layer which causes mobileBERT model to seg fault
when running in concurrent mode.

Signed-off-by: mathad01 <matthew.haddon@arm.com>
Change-Id: Ib2b73ffaa5b70147b837b7c764064e530baba637
---
 src/backends/reference/workloads/RefStackWorkload.cpp | 5 ++++-
 src/backends/reference/workloads/Stack.cpp            | 7 +++----
 src/backends/reference/workloads/Stack.hpp            | 4 +++-
 3 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/src/backends/reference/workloads/RefStackWorkload.cpp b/src/backends/reference/workloads/RefStackWorkload.cpp
index 31949e967e..3f7fd7bda2 100644
--- a/src/backends/reference/workloads/RefStackWorkload.cpp
+++ b/src/backends/reference/workloads/RefStackWorkload.cpp
@@ -32,6 +32,9 @@ void RefStackWorkload::Execute(std::vector<ITensorHandle*> inputs, std::vector<I
 {
     ARMNN_SCOPED_PROFILING_EVENT(Compute::CpuRef, "RefStackWorkload_Execute");
 
+    const TensorInfo& inputInfo  = GetTensorInfo(inputs[0]);
+    const TensorInfo& outputInfo = GetTensorInfo(outputs[0]);
+
     std::vector<std::unique_ptr<Decoder<float>>> inputDecoders;
     for (unsigned int i=0; i<inputs.size(); ++i)
     {
@@ -41,7 +44,7 @@ void RefStackWorkload::Execute(std::vector<ITensorHandle*> inputs, std::vector<I
     std::unique_ptr<Encoder<float>> outputEncoder = MakeEncoder<float>(GetTensorInfo(outputs[0]),
                                                                        outputs[0]->Map());
 
-    Stack(m_Data, inputDecoders, *outputEncoder);
+    Stack(m_Data, inputDecoders, *outputEncoder, inputInfo, outputInfo);
 }
 
 } // namespace armnn
diff --git a/src/backends/reference/workloads/Stack.cpp b/src/backends/reference/workloads/Stack.cpp
index f2bce54d6a..ca8b51a98d 100644
--- a/src/backends/reference/workloads/Stack.cpp
+++ b/src/backends/reference/workloads/Stack.cpp
@@ -11,11 +11,10 @@ namespace armnn
 
 void Stack(const StackQueueDescriptor& data,
            std::vector<std::unique_ptr<Decoder<float>>>& inputs,
-           Encoder<float>& output)
+           Encoder<float>& output,
+           const TensorInfo& inputInfo,
+           const TensorInfo& outputInfo)
 {
-    const TensorInfo& outputInfo = GetTensorInfo(data.m_Outputs[0]);
-    const TensorInfo& inputInfo = GetTensorInfo(data.m_Inputs[0]);
-
     unsigned int outputNumDims = outputInfo.GetNumDimensions();
     unsigned int inputNumDims = inputInfo.GetNumDimensions();
 
diff --git a/src/backends/reference/workloads/Stack.hpp b/src/backends/reference/workloads/Stack.hpp
index cd86d41552..2f02a4c644 100644
--- a/src/backends/reference/workloads/Stack.hpp
+++ b/src/backends/reference/workloads/Stack.hpp
@@ -15,6 +15,8 @@ namespace armnn
 
 void Stack (const StackQueueDescriptor&                   data,
             std::vector<std::unique_ptr<Decoder<float>>>& inputs,
-            Encoder<float>&                               output);
+            Encoder<float>&                               output,
+            const TensorInfo&                             inputInfo,
+            const TensorInfo&                             outputInfo);
 
 } // namespace armnn
-- 
cgit v1.2.1