1 files changed, 32 insertions, 18 deletions

diff --git a/src/backends/reference/workloads/RefConvolution2dWorkload.cpp b/src/backends/reference/workloads/RefConvolution2dWorkload.cpp
index d57040eaec..fe97cb1066 100644
--- a/src/backends/reference/workloads/RefConvolution2dWorkload.cpp
+++ b/src/backends/reference/workloads/RefConvolution2dWorkload.cpp
@@ -12,37 +12,46 @@
 
 namespace armnn
 {
-RefConvolution2dWorkload::RefConvolution2dWorkload(
-    const Convolution2dQueueDescriptor& descriptor, const WorkloadInfo& info)
+RefConvolution2dWorkload::RefConvolution2dWorkload(const Convolution2dQueueDescriptor& descriptor,
+                                                   const WorkloadInfo& info)
     : RefBaseWorkload<Convolution2dQueueDescriptor>(descriptor, info)
 {
     WorkloadInfo detailsInfo;
     detailsInfo.m_InputTensorInfos = info.m_InputTensorInfos;
     detailsInfo.m_OutputTensorInfos = info.m_OutputTensorInfos;
-    detailsInfo.m_WeightsTensorInfo = armnn::Optional<armnn::TensorInfo>(descriptor.m_Weight->GetTensorInfo());
-    if (descriptor.m_Parameters.m_BiasEnabled)
-    {
-        detailsInfo.m_BiasTensorInfo = armnn::Optional<armnn::TensorInfo>(descriptor.m_Bias->GetTensorInfo());
-    }
 
     // Report Profiling Details
     ARMNN_REPORT_PROFILING_WORKLOAD_DESC("RefConvolution2dWorkload_Construct",
                                          descriptor.m_Parameters,
                                          detailsInfo,
                                          this->GetGuid());
+}
 
-    m_Weight = std::make_unique<ScopedTensorHandle>(*( descriptor.m_Weight ));
-    const TensorInfo& rFilterInfo = m_Weight->GetTensorInfo();
+void RefConvolution2dWorkload::PostAllocationConfigure()
+{
+    PostAllocationConfigure(m_Data.m_Inputs, m_Data.m_Outputs);
+}
 
+void RefConvolution2dWorkload::PostAllocationConfigure(std::vector<ITensorHandle*> inputs,
+                                                        std::vector<ITensorHandle*> outputs)
+{
+    const TensorInfo& inputInfo = GetTensorInfo(inputs[0]);
+    ARMNN_ASSERT(inputInfo.GetNumDimensions() > 1);
+    m_InputShape = inputInfo.GetShape();
+
+    const TensorInfo& rFilterInfo = GetTensorInfo(inputs[1]);
+    ARMNN_ASSERT(inputInfo.GetNumDimensions() > 1);
     m_FilterShape = rFilterInfo.GetShape();
-    m_FilterDecoder = MakeDecoder<float>(rFilterInfo, m_Weight.get()->Map(true));
+    m_FilterDecoder = MakeDecoder<float>(rFilterInfo);
 
-    if ( descriptor.m_Parameters.m_BiasEnabled )
+    if (m_Data.m_Parameters.m_BiasEnabled)
     {
-        m_Bias = std::make_unique<ScopedTensorHandle>(*( descriptor.m_Bias ));
-        const TensorInfo& biasInfo = m_Bias->GetTensorInfo();
-        m_BiasDecoder = MakeDecoder<float>(biasInfo, m_Bias->Map(true));
+        const TensorInfo& biasInfo = GetTensorInfo(inputs[2]);
+        m_BiasDecoder = MakeDecoder<float>(biasInfo);
     }
+
+    const TensorInfo& outputInfo = GetTensorInfo(outputs[0]);
+    m_OutputShape = outputInfo.GetShape();
 }
 
 void RefConvolution2dWorkload::Execute() const
@@ -52,6 +61,8 @@ void RefConvolution2dWorkload::Execute() const
 
 void RefConvolution2dWorkload::ExecuteAsync(WorkingMemDescriptor& workingMemDescriptor)
 {
+    PostAllocationConfigure(workingMemDescriptor.m_Inputs, workingMemDescriptor.m_Outputs);
+
     Execute(workingMemDescriptor.m_Inputs, workingMemDescriptor.m_Outputs);
 }
 
@@ -62,14 +73,17 @@ void RefConvolution2dWorkload::Execute(std::vector<ITensorHandle*> inputs, std::
     std::unique_ptr<Decoder<float>> inputDecoder = MakeDecoder<float>(GetTensorInfo(inputs[0]), inputs[0]->Map());
     std::unique_ptr<Encoder<float>> outputEncoder = MakeEncoder<float>(GetTensorInfo(outputs[0]), outputs[0]->Map());
 
-    const TensorShape& inputShape = GetTensorInfo(inputs[0]).GetShape();
-    const TensorShape& outputShape = GetTensorInfo(outputs[0]).GetShape();
+    m_FilterDecoder->Reset(inputs[1]->Map());
+    if (m_Data.m_Parameters.m_BiasEnabled)
+    {
+        m_BiasDecoder->Reset(inputs[2]->Map());
+    }
 
-    Convolve(inputShape, *inputDecoder, outputShape, *outputEncoder, m_FilterShape,
+    Convolve(m_InputShape, *inputDecoder, m_OutputShape, *outputEncoder, m_FilterShape,
              *m_FilterDecoder, m_Data.m_Parameters.m_BiasEnabled, m_BiasDecoder.get(),
              m_Data.m_Parameters.m_DataLayout, m_Data.m_Parameters.m_PadTop, m_Data.m_Parameters.m_PadLeft,
              m_Data.m_Parameters.m_StrideX, m_Data.m_Parameters.m_StrideY,
              m_Data.m_Parameters.m_DilationX, m_Data.m_Parameters.m_DilationY);
 }
 
-} //namespace armnn
+} //namespace armnn
\ No newline at end of file