IVGCVSW-2866 Implement RegisterDebugCallback for RefDebugWorkload

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

7 files changed, 132 insertions, 15 deletions

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 3f087dbce9..ec237aab04 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -409,6 +409,7 @@ if(BUILD_UNIT_TESTS)
         src/armnn/test/ConstTensorLayerVisitor.cpp
         src/armnn/test/CreateWorkload.hpp
         src/armnn/test/CsvReaderTest.cpp
+        src/armnn/test/DebugCallbackTest.cpp
         src/armnn/test/EndToEndTest.cpp
         src/armnn/test/FloatingPointConverterTest.cpp
         src/armnn/test/GraphTests.cpp
diff --git a/include/armnn/Types.hpp b/include/armnn/Types.hpp
index bdddd37a1c..438b6be741 100644
--- a/include/armnn/Types.hpp
+++ b/include/armnn/Types.hpp
@@ -183,7 +183,10 @@ using LayerGuid = unsigned int;
 
 class ITensorHandle;
 
-/// Define the type of callback for the debug layer to call
-using DebugCallbackFunction = std::function<void(LayerGuid, unsigned int, ITensorHandle*)>;
+/// Define the type of callback for the Debug layer to call
+/// @param guid - guid of layer connected to the input of the Debug layer
+/// @param slotIndex - index of the output slot connected to the input of the Debug layer
+/// @param tensorHandle - TensorHandle for the input tensor to the Debug layer
+using DebugCallbackFunction = std::function<void(LayerGuid guid, unsigned int slotIndex, ITensorHandle* tensorHandle)>;
 
 } // namespace armnn
diff --git a/src/armnn/test/DebugCallbackTest.cpp b/src/armnn/test/DebugCallbackTest.cpp
new file mode 100644
index 0000000000..0b1abbf2dc
--- /dev/null
+++ b/src/armnn/test/DebugCallbackTest.cpp
@@ -0,0 +1,98 @@
+//
+// Copyright © 2017 Arm Ltd. All rights reserved.
+// SPDX-License-Identifier: MIT
+//
+
+#include <armnn/Descriptors.hpp>
+#include <armnn/IRuntime.hpp>
+#include <armnn/INetwork.hpp>
+#include <armnn/Types.hpp>
+#include <Runtime.hpp>
+
+#include <boost/test/unit_test.hpp>
+
+BOOST_AUTO_TEST_SUITE(DebugCallback)
+
+namespace
+{
+
+using namespace armnn;
+
+INetworkPtr CreateSimpleNetwork()
+{
+    INetworkPtr net(INetwork::Create());
+
+    IConnectableLayer* input = net->AddInputLayer(0, "Input");
+
+    ActivationDescriptor descriptor;
+    descriptor.m_Function = ActivationFunction::ReLu;
+    IConnectableLayer* activationLayer = net->AddActivationLayer(descriptor, "Activation:ReLu");
+
+    IConnectableLayer* output = net->AddOutputLayer(0);
+
+    input->GetOutputSlot(0).Connect(activationLayer->GetInputSlot(0));
+    activationLayer->GetOutputSlot(0).Connect(output->GetInputSlot(0));
+
+    input->GetOutputSlot(0).SetTensorInfo(TensorInfo({ 1, 1, 1, 5 }, DataType::Float32));
+    activationLayer->GetOutputSlot(0).SetTensorInfo(TensorInfo({ 1, 1, 1, 5 }, DataType::Float32));
+
+    return net;
+}
+
+BOOST_AUTO_TEST_CASE(RuntimeRegisterDebugCallback)
+{
+    INetworkPtr net = CreateSimpleNetwork();
+
+    IRuntime::CreationOptions options;
+    IRuntimePtr runtime(IRuntime::Create(options));
+
+    // Optimize the network with debug option
+    OptimizerOptions optimizerOptions(false, true);
+    std::vector<BackendId> backends = { "CpuRef" };
+    IOptimizedNetworkPtr optNet = Optimize(*net, backends, runtime->GetDeviceSpec(), optimizerOptions);
+
+    NetworkId netId;
+    BOOST_TEST(runtime->LoadNetwork(netId, std::move(optNet)) == Status::Success);
+
+    // Set up callback function
+    int callCount = 0;
+    std::vector<TensorShape> tensorShapes;
+    std::vector<unsigned int> slotIndexes;
+    auto mockCallback = [&](LayerGuid guid, unsigned int slotIndex, ITensorHandle* tensor)
+    {
+        slotIndexes.push_back(slotIndex);
+        tensorShapes.push_back(tensor->GetShape());
+        callCount++;
+    };
+
+    runtime->RegisterDebugCallback(netId, mockCallback);
+
+    std::vector<float> inputData({-2, -1, 0, 1, 2});
+    std::vector<float> outputData(5);
+
+    InputTensors inputTensors
+    {
+        {0, ConstTensor(runtime->GetInputTensorInfo(netId, 0), inputData.data())}
+    };
+    OutputTensors outputTensors
+    {
+        {0, Tensor(runtime->GetOutputTensorInfo(netId, 0), outputData.data())}
+    };
+
+    runtime->EnqueueWorkload(netId, inputTensors, outputTensors);
+
+    // Check that the callback was called twice
+    BOOST_TEST(callCount == 2);
+
+    // Check that tensor handles passed to callback have correct shapes
+    const std::vector<TensorShape> expectedShapes({TensorShape({1, 1, 1, 5}), TensorShape({1, 1, 1, 5})});
+    BOOST_TEST(tensorShapes == expectedShapes);
+
+    // Check that slot indexes passed to callback are correct
+    const std::vector<unsigned int> expectedSlotIndexes({0, 0});
+    BOOST_TEST(slotIndexes == expectedSlotIndexes);
+}
+
+} // anonymous namespace
+
+BOOST_AUTO_TEST_SUITE_END()
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>;