commit 788e2a6c917abe7c5187d9e5c349683d456080e5
author Teresa Charlin <teresa.charlinreyes@arm.com> Mon Jan 17 21:19:52 2022 +0000
committer Teresa Charlin <teresa.charlinreyes@arm.com> Tue Jan 18 21:32:27 2022 +0000
tree 1c88b36828a3ebb6ef04c3e8c0bc92bdd330e9cf
parent adeebaa73205bd981ea7e8c8f135f01355cba841

IVGCVSW-6682 Add ReplaceTensorHandle functions to IWorkload and BaseWorkload

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

include/armnn/backends/IWorkload.hpp

diff --git a/include/armnn/backends/IWorkload.hpp b/include/armnn/backends/IWorkload.hpp
index a4827eb..d63e0ac 100644
--- a/include/armnn/backends/IWorkload.hpp
+++ b/include/armnn/backends/IWorkload.hpp
@@ -31,6 +31,12 @@
 
     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*/) {}
 };
 

include/armnn/backends/Workload.hpp

diff --git a/include/armnn/backends/Workload.hpp b/include/armnn/backends/Workload.hpp
index 7c1bda5..07e1abb 100644
--- a/include/armnn/backends/Workload.hpp
+++ b/include/armnn/backends/Workload.hpp
@@ -54,6 +54,18 @@
 
     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;