IVGCVSW-3215 Add CreateWorkload test for Rsqrt Fp32

Change-Id: Iea6edf90148773b12f361bfa4b1d572c5322d18b
Signed-off-by: nikraj01 <nikhil.raj@arm.com>

2 files changed, 48 insertions, 0 deletions

diff --git a/src/armnn/test/CreateWorkload.hpp b/src/armnn/test/CreateWorkload.hpp
index 135a4421cd..87df00af3c 100644
--- a/src/armnn/test/CreateWorkload.hpp
+++ b/src/armnn/test/CreateWorkload.hpp
@@ -892,6 +892,34 @@ std::unique_ptr<ResizeBilinearWorkload> CreateResizeBilinearWorkloadTest(armnn::
     return workload;
 }
 
+template <typename RsqrtWorkload, armnn::DataType DataType>
+std::unique_ptr<RsqrtWorkload> CreateRsqrtWorkloadTest(armnn::IWorkloadFactory& factory,
+                                                       armnn::Graph&  graph)
+{
+    Layer* const layer = graph.AddLayer<RsqrtLayer>("rsqrt");
+
+    // Creates extra layers.
+    Layer* const input = graph.AddLayer<InputLayer>(0, "input");
+    Layer* const output = graph.AddLayer<OutputLayer>(0, "output");
+
+    // Connects up.
+    armnn::TensorInfo tensorInfo({1, 1}, DataType);
+
+    Connect(input, layer, tensorInfo);
+    Connect(layer, output, tensorInfo);
+
+    CreateTensorHandles(graph, factory);
+
+    // Makes the workload and checks it.
+    auto workload = MakeAndCheckWorkload<RsqrtWorkload>(*layer, graph, factory);
+
+    RsqrtQueueDescriptor queueDescriptor = workload->GetData();
+    BOOST_TEST(queueDescriptor.m_Inputs.size() == 1);
+    BOOST_TEST(queueDescriptor.m_Outputs.size() == 1);
+
+    return workload;
+}
+
 template <typename L2NormalizationWorkload, armnn::DataType DataType>
 std::unique_ptr<L2NormalizationWorkload> CreateL2NormalizationWorkloadTest(armnn::IWorkloadFactory& factory,
     armnn::Graph& graph, DataLayout dataLayout = DataLayout::NCHW)
diff --git a/src/backends/reference/test/RefCreateWorkloadTests.cpp b/src/backends/reference/test/RefCreateWorkloadTests.cpp
index 0311276f10..7c5712b915 100644
--- a/src/backends/reference/test/RefCreateWorkloadTests.cpp
+++ b/src/backends/reference/test/RefCreateWorkloadTests.cpp
@@ -614,6 +614,26 @@ BOOST_AUTO_TEST_CASE(CreateResizeBilinearFloat32Nhwc)
     RefCreateResizeBilinearTest<RefResizeBilinearFloat32Workload, armnn::DataType::Float32>(DataLayout::NHWC);
 }
 
+template <typename RsqrtWorkloadType, armnn::DataType DataType>
+static void RefCreateRsqrtTest()
+{
+    Graph graph;
+    RefWorkloadFactory factory;
+
+    auto workload = CreateRsqrtWorkloadTest<RsqrtWorkloadType, DataType>(factory, graph);
+
+    // Checks that outputs are as we expect them (see definition of CreateRsqrtWorkloadTest).
+    CheckInputOutput(std::move(workload),
+                     TensorInfo({ 1, 1 }, DataType),
+                     TensorInfo({ 1, 1 }, DataType));
+
+}
+
+BOOST_AUTO_TEST_CASE(CreateRsqrtFloat32)
+{
+    RefCreateRsqrtTest<RefRsqrtFloat32Workload, armnn::DataType::Float32>();
+}
+
 template <typename L2NormalizationWorkloadType, armnn::DataType DataType>
 static void RefCreateL2NormalizationTest(DataLayout dataLayout)
 {