IVGCVSW-6682 Add ReplaceTensorHandle functions to IWorkload and BaseWorkload

Signed-off-by: Teresa Charlin <teresa.charlinreyes@arm.com>
Change-Id: I9f80b9f45206db920568e28e363fcb60f5c0819a

2 files changed, 18 insertions, 0 deletions

diff --git a/include/armnn/backends/IWorkload.hpp b/include/armnn/backends/IWorkload.hpp
index a4827ebcdf..d63e0acc72 100644
--- a/include/armnn/backends/IWorkload.hpp
+++ b/include/armnn/backends/IWorkload.hpp
@@ -31,6 +31,12 @@ public:
 
     virtual profiling::ProfilingGuid GetGuid() const = 0;
 
+    // Replace input tensor handle with the given TensorHandle
+    virtual void ReplaceInputTensorHandle(ITensorHandle* /*input*/, unsigned int /*slot*/) = 0;
+
+    // Replace output tensor handle with the given TensorHandle
+    virtual void ReplaceOutputTensorHandle(ITensorHandle* /*output*/, unsigned int /*slot*/) = 0;
+
     virtual void RegisterDebugCallback(const DebugCallbackFunction& /*func*/) {}
 };
 
diff --git a/include/armnn/backends/Workload.hpp b/include/armnn/backends/Workload.hpp
index 7c1bda50bc..07e1abb392 100644
--- a/include/armnn/backends/Workload.hpp
+++ b/include/armnn/backends/Workload.hpp
@@ -54,6 +54,18 @@ public:
 
     profiling::ProfilingGuid GetGuid() const final { return m_Guid; }
 
+    // Replace input tensor handle with the given TensorHandle
+    void ReplaceInputTensorHandle(ITensorHandle* tensorHandle, unsigned int slot) override
+    {
+        m_Data.m_Inputs[slot] = tensorHandle;
+    }
+
+    // Replace output tensor handle with the given TensorHandle
+    void ReplaceOutputTensorHandle(ITensorHandle* tensorHandle, unsigned int slot) override
+    {
+        m_Data.m_Outputs[slot] = tensorHandle;
+    }
+
 protected:
     QueueDescriptor m_Data;
     const profiling::ProfilingGuid m_Guid;