From 788e2a6c917abe7c5187d9e5c349683d456080e5 Mon Sep 17 00:00:00 2001 From: Teresa Charlin Date: Mon, 17 Jan 2022 21:19:52 +0000 Subject: IVGCVSW-6682 Add ReplaceTensorHandle functions to IWorkload and BaseWorkload Signed-off-by: Teresa Charlin Change-Id: I9f80b9f45206db920568e28e363fcb60f5c0819a --- include/armnn/backends/IWorkload.hpp | 6 ++++++ include/armnn/backends/Workload.hpp | 12 ++++++++++++ 2 files changed, 18 insertions(+) (limited to 'include/armnn') 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; -- cgit v1.2.1