IVGCVSW-2866 Implement RegisterDebugCallback for RefDebugWorkload

Change-Id: I9144fb6b7d05561b5b8fd9db5dbe31c9257f10ca
Signed-off-by: Nattapat Chaimanowong <nattapat.chaimanowong@arm.com>

4 files changed, 28 insertions, 13 deletions

diff --git a/src/backends/reference/workloads/Debug.cpp b/src/backends/reference/workloads/Debug.cpp
index b263db67cf..594f428908 100644
--- a/src/backends/reference/workloads/Debug.cpp
+++ b/src/backends/reference/workloads/Debug.cpp
@@ -6,7 +6,6 @@
 
 #include <boost/numeric/conversion/cast.hpp>
 
-#include <cstring>
 #include <algorithm>
 #include <iostream>
 
@@ -15,9 +14,7 @@ namespace armnn
 
 template <typename T>
 void Debug(const TensorInfo& inputInfo,
-           const TensorInfo& outputInfo,
            const T* inputData,
-           T* outputData,
            LayerGuid guid,
            const std::string& layerName,
            unsigned int slotIndex)
@@ -86,22 +83,16 @@ void Debug(const TensorInfo& inputInfo,
     }
 
     std::cout << " }" << std::endl;
-
-    std::memcpy(outputData, inputData, inputInfo.GetNumElements()*sizeof(T));
 }
 
 template void Debug<float>(const TensorInfo& inputInfo,
-                           const TensorInfo& outputInfo,
                            const float* inputData,
-                           float* outputData,
                            LayerGuid guid,
                            const std::string& layerName,
                            unsigned int slotIndex);
 
 template void Debug<uint8_t>(const TensorInfo& inputInfo,
-                             const TensorInfo& outputInfo,
                              const uint8_t* inputData,
-                             uint8_t* outputData,
                              LayerGuid guid,
                              const std::string& layerName,
                              unsigned int slotIndex);
diff --git a/src/backends/reference/workloads/Debug.hpp b/src/backends/reference/workloads/Debug.hpp
index 29a7d40662..3f9920c543 100644
--- a/src/backends/reference/workloads/Debug.hpp
+++ b/src/backends/reference/workloads/Debug.hpp
@@ -11,9 +11,7 @@ namespace armnn
 
 template <typename T>
 void Debug(const TensorInfo& inputInfo,
-           const TensorInfo& outputInfo,
            const T* inputData,
-           T* outputData,
            LayerGuid guid,
            const std::string& layerName,
            unsigned int slotIndex);
diff --git a/src/backends/reference/workloads/RefDebugWorkload.cpp b/src/backends/reference/workloads/RefDebugWorkload.cpp
index 412d399adc..3da925913a 100644
--- a/src/backends/reference/workloads/RefDebugWorkload.cpp
+++ b/src/backends/reference/workloads/RefDebugWorkload.cpp
@@ -9,6 +9,8 @@
 
 #include <TypeUtils.hpp>
 
+#include <cstring>
+
 namespace armnn
 {
 
@@ -20,12 +22,27 @@ void RefDebugWorkload<DataType>::Execute() const
     ARMNN_SCOPED_PROFILING_EVENT(Compute::CpuRef, GetName() + "_Execute");
 
     const TensorInfo& inputInfo = GetTensorInfo(m_Data.m_Inputs[0]);
-    const TensorInfo& outputInfo = GetTensorInfo(m_Data.m_Outputs[0]);
 
     const T* inputData = GetInputTensorData<T>(0, m_Data);
     T* outputData = GetOutputTensorData<T>(0, m_Data);
 
-    Debug(inputInfo, outputInfo, inputData, outputData, m_Data.m_Guid, m_Data.m_LayerName, m_Data.m_SlotIndex);
+    if (m_Callback)
+    {
+        m_Callback(m_Data.m_Guid, m_Data.m_SlotIndex, m_Data.m_Inputs[0]);
+    }
+    else
+    {
+        Debug(inputInfo, inputData, m_Data.m_Guid, m_Data.m_LayerName, m_Data.m_SlotIndex);
+    }
+
+    std::memcpy(outputData, inputData, inputInfo.GetNumElements()*sizeof(T));
+
+}
+
+template<armnn::DataType DataType>
+void RefDebugWorkload<DataType>::RegisterDebugCallback(const DebugCallbackFunction& func)
+{
+    m_Callback = func;
 }
 
 template class RefDebugWorkload<DataType::Float32>;
diff --git a/src/backends/reference/workloads/RefDebugWorkload.hpp b/src/backends/reference/workloads/RefDebugWorkload.hpp
index c1a3e26ec2..2985699f7b 100644
--- a/src/backends/reference/workloads/RefDebugWorkload.hpp
+++ b/src/backends/reference/workloads/RefDebugWorkload.hpp
@@ -16,6 +16,10 @@ template <armnn::DataType DataType>
 class RefDebugWorkload : public TypedWorkload<DebugQueueDescriptor, DataType>
 {
 public:
+    RefDebugWorkload(const DebugQueueDescriptor& descriptor, const WorkloadInfo& info)
+    : TypedWorkload<DebugQueueDescriptor, DataType>(descriptor, info)
+    , m_Callback(nullptr) {}
+
     static const std::string& GetName()
     {
         static const std::string name = std::string("RefDebug") + GetDataTypeName(DataType) + "Workload";
@@ -26,6 +30,11 @@ public:
     using TypedWorkload<DebugQueueDescriptor, DataType>::TypedWorkload;
 
     void Execute() const override;
+
+    void RegisterDebugCallback(const DebugCallbackFunction& func) override;
+
+private:
+    DebugCallbackFunction m_Callback;
 };
 
 using RefDebugFloat32Workload = RefDebugWorkload<DataType::Float32>;