From 6e9482013f41725ccca0767c0c5db9b29f77d981 Mon Sep 17 00:00:00 2001
From: Nattapat Chaimanowong <nattapat.chaimanowong@arm.com>
Date: Fri, 22 Mar 2019 14:01:46 +0000
Subject: IVGCVSW-2865 Extend IRuntime to add a new method
 RegisterDebugCallback(...)

* Made changes to LoadedNetwork and IWorkload to pass on the registered
callback function

Change-Id: I6ea10f2a299d6de8bf681c8ff36d3fbed1d6d887
Signed-off-by: Nattapat Chaimanowong <nattapat.chaimanowong@arm.com>
---
 src/armnn/LoadedNetwork.cpp | 8 ++++++++
 src/armnn/LoadedNetwork.hpp | 2 ++
 src/armnn/Runtime.cpp       | 6 ++++++
 src/armnn/Runtime.hpp       | 5 +++++
 4 files changed, 21 insertions(+)

(limited to 'src/armnn')

diff --git a/src/armnn/LoadedNetwork.cpp b/src/armnn/LoadedNetwork.cpp
index 4221d36036..9263f1a6e9 100644
--- a/src/armnn/LoadedNetwork.cpp
+++ b/src/armnn/LoadedNetwork.cpp
@@ -485,4 +485,12 @@ bool LoadedNetwork::Execute()
     return success;
 }
 
+void LoadedNetwork::RegisterDebugCallback(const DebugCallbackFunction& func)
+{
+    for (auto&& workloadPtr: m_WorkloadQueue)
+    {
+        workloadPtr.get()->RegisterDebugCallback(func);
+    }
+}
+
 }
diff --git a/src/armnn/LoadedNetwork.hpp b/src/armnn/LoadedNetwork.hpp
index 9c0fe0b108..75af4a4e28 100644
--- a/src/armnn/LoadedNetwork.hpp
+++ b/src/armnn/LoadedNetwork.hpp
@@ -49,6 +49,8 @@ public:
 
     void FreeWorkingMemory();
 
+    void RegisterDebugCallback(const DebugCallbackFunction& func);
+
 private:
     void AllocateWorkingMemory();
 
diff --git a/src/armnn/Runtime.cpp b/src/armnn/Runtime.cpp
index 09be92c709..f8b2462f96 100644
--- a/src/armnn/Runtime.cpp
+++ b/src/armnn/Runtime.cpp
@@ -231,4 +231,10 @@ Status Runtime::EnqueueWorkload(NetworkId networkId,
     return loadedNetwork->EnqueueWorkload(inputTensors, outputTensors);
 }
 
+void Runtime::RegisterDebugCallback(NetworkId networkId, const DebugCallbackFunction& func)
+{
+    LoadedNetwork* loadedNetwork = GetLoadedNetworkPtr(networkId);
+    loadedNetwork->RegisterDebugCallback(func);
+}
+
 }
diff --git a/src/armnn/Runtime.hpp b/src/armnn/Runtime.hpp
index a3f4a3930b..10383bc970 100644
--- a/src/armnn/Runtime.hpp
+++ b/src/armnn/Runtime.hpp
@@ -59,6 +59,11 @@ public:
     /// @return A pointer to the requested profiler, or nullptr if not found.
     virtual const std::shared_ptr<IProfiler> GetProfiler(NetworkId networkId) const override;
 
+    /// Registers a callback function to debug layers performing custom computations on intermediate tensors.
+    /// @param networkId The id of the network to register the callback.
+    /// @param func callback function to pass to the debug layer.
+    virtual void RegisterDebugCallback(NetworkId networkId, const DebugCallbackFunction& func) override;
+
     /// Creates a runtime for workload execution.
     /// May throw a ClRuntimeUnavailableException if @a defaultComputeDevice requires a CL runtime but
     /// it cannot be setup for some reason.
-- 
cgit v1.2.1