3 files changed, 46 insertions, 4 deletions
diff --git a/src/backends/cl/workloads/ClBaseWorkload.hpp b/src/backends/cl/workloads/ClBaseWorkload.hpp
index e74fc84f4f..03417e33ae 100644
--- a/src/backends/cl/workloads/ClBaseWorkload.hpp
+++ b/src/backends/cl/workloads/ClBaseWorkload.hpp
@@ -20,17 +20,38 @@ public:
     // Replace input tensor handle with the given TensorHandle and call Reconfigure()
     void ReplaceInputTensorHandle(ITensorHandle* tensorHandle, unsigned int slot) override
     {
+        ITensorHandle* backupHandle = this->m_Data.m_Inputs[slot];
         this->m_Data.m_Inputs[slot] = tensorHandle;
-        Reconfigure();
+        try
+        {
+            Reconfigure();
+        }
+        catch(armnn::UnimplementedException& e)
+        {
+            // Cannot reconfigure, revert the slot back and throw the exception.
+            this->m_Data.m_Inputs[slot] = backupHandle;
+            throw e;
+        }
     }
 
     // Replace output tensor handle with the given TensorHandle and call Reconfigure()
     void ReplaceOutputTensorHandle(ITensorHandle* tensorHandle, unsigned int slot) override
     {
+        ITensorHandle* backupHandle = this->m_Data.m_Outputs[slot];
         this->m_Data.m_Outputs[slot] = tensorHandle;
-        Reconfigure();
+        try
+        {
+            Reconfigure();
+        }
+        catch(armnn::UnimplementedException& e)
+        {
+            // Cannot reconfigure, revert the slot back and throw the exception.
+            this->m_Data.m_Inputs[slot] = backupHandle;
+            throw e;
+        }
     }
 
+protected:
     // Reconfigure the workload configuration. Throw armnn::UnimplementedException by default.
     virtual void Reconfigure()
     {
diff --git a/src/backends/cl/workloads/ClConvolution2dWorkload.cpp b/src/backends/cl/workloads/ClConvolution2dWorkload.cpp
index 705e92d307..cdfa885f67 100644
--- a/src/backends/cl/workloads/ClConvolution2dWorkload.cpp
+++ b/src/backends/cl/workloads/ClConvolution2dWorkload.cpp
@@ -90,6 +90,10 @@ ClConvolution2dWorkload::ClConvolution2dWorkload(const Convolution2dQueueDescrip
     arm_compute::ICLTensor& input  = static_cast<IClTensorHandle*>(m_Data.m_Inputs[0])->GetTensor();
     arm_compute::ICLTensor& output = static_cast<IClTensorHandle*>(m_Data.m_Outputs[0])->GetTensor();
 
+    // Create Proxy tensor and set the initial tensor handle to it
+    m_InputProxy = std::make_unique<ICLTensorProxy>(&input);
+    m_OutputProxy = std::make_unique<ICLTensorProxy>(&output);
+
     arm_compute::DataLayout aclDataLayout = ConvertDataLayout(m_Data.m_Parameters.m_DataLayout);
     input.info()->set_data_layout(aclDataLayout);
     output.info()->set_data_layout(aclDataLayout);
@@ -101,10 +105,10 @@ ClConvolution2dWorkload::ClConvolution2dWorkload(const Convolution2dQueueDescrip
     {
         ARMNN_SCOPED_PROFILING_EVENT(Compute::Undefined, "ClConvolution2dWorkload_configure");
         m_ConvolutionLayer.configure(clCompileContext,
-                                     &input,
+                                     m_InputProxy.get(),
                                      m_KernelTensor.get(),
                                      m_BiasTensor.get(),
-                                     &output,
+                                     m_OutputProxy.get(),
                                      padStrideInfo,
                                      arm_compute::WeightsInfo(),
                                      aclDilationInfo,
@@ -174,4 +178,13 @@ void ClConvolution2dWorkload::FreeUnusedTensors()
     FreeTensorIfUnused(m_BiasTensor);
 }
 
+void ClConvolution2dWorkload::Reconfigure()
+{
+    ARMNN_SCOPED_PROFILING_EVENT(Compute::Undefined, "ClConvolution2dWorkload_Reconfigure");
+    arm_compute::ICLTensor& input  = static_cast<IClTensorHandle*>(m_Data.m_Inputs[0])->GetTensor();
+    arm_compute::ICLTensor& output = static_cast<IClTensorHandle*>(m_Data.m_Outputs[0])->GetTensor();
+    m_InputProxy->set(&input);
+    m_OutputProxy->set(&output);
+}
+
 } //namespace armnn
diff --git a/src/backends/cl/workloads/ClConvolution2dWorkload.hpp b/src/backends/cl/workloads/ClConvolution2dWorkload.hpp
index 8a4599df47..891d5096cd 100644
--- a/src/backends/cl/workloads/ClConvolution2dWorkload.hpp
+++ b/src/backends/cl/workloads/ClConvolution2dWorkload.hpp
@@ -13,6 +13,8 @@
 #include <arm_compute/runtime/CL/functions/CLConvolutionLayer.h>
 #include <arm_compute/runtime/MemoryManagerOnDemand.h>
 
+#include <cl/ICLTensorProxy.hpp>
+
 #include <memory>
 
 namespace armnn
@@ -38,6 +40,9 @@ public:
 
     arm_compute::ConvolutionMethod GetConvolutionMethod() const;
 
+protected:
+    void Reconfigure() override;
+
 private:
     mutable arm_compute::CLConvolutionLayer m_ConvolutionLayer;
 
@@ -47,6 +52,9 @@ private:
     arm_compute::ConvolutionMethod m_ConvolutionMethod;
 
     void FreeUnusedTensors();
+
+    std::unique_ptr<ICLTensorProxy> m_InputProxy;
+    std::unique_ptr<ICLTensorProxy> m_OutputProxy;
 };
 
 } //namespace armnn