From a3b20473ea0583ba4c6a2c17ba2d59ed958c0aa7 Mon Sep 17 00:00:00 2001
From: Teresa Charlin <teresa.charlinreyes@arm.com>
Date: Thu, 6 Jun 2019 11:12:32 +0100
Subject: IVGCVSW-3142 Refactor reference Pooling2d workload

Signed-off-by: Teresa Charlin <teresa.charlinreyes@arm.com>
Change-Id: I94c973ab747309c0214268c9c39f6d8f3fc7b255
---
 src/backends/reference/RefLayerSupport.cpp | 24 +++++++++++++++++++-----
 1 file changed, 19 insertions(+), 5 deletions(-)

(limited to 'src/backends/reference/RefLayerSupport.cpp')

diff --git a/src/backends/reference/RefLayerSupport.cpp b/src/backends/reference/RefLayerSupport.cpp
index adc63e965b..f177385837 100644
--- a/src/backends/reference/RefLayerSupport.cpp
+++ b/src/backends/reference/RefLayerSupport.cpp
@@ -1004,12 +1004,26 @@ bool RefLayerSupport::IsPooling2dSupported(const TensorInfo& input,
                                            const Pooling2dDescriptor& descriptor,
                                            Optional<std::string&> reasonIfUnsupported) const
 {
-    ignore_unused(output);
     ignore_unused(descriptor);
-    return IsSupportedForDataTypeRef(reasonIfUnsupported,
-                                     input.GetDataType(),
-                                     &TrueFunc<>,
-                                     &TrueFunc<>);
+    bool supported = true;
+
+    // Define supported output and inputs types.
+    std::array<DataType,2> supportedTypes =
+    {
+        DataType::Float32,
+        DataType::QuantisedAsymm8
+    };
+
+    supported &= CheckSupportRule(TypeAnyOf(input, supportedTypes), reasonIfUnsupported,
+                                  "Reference poolind2d: input is not a supported type.");
+
+    supported &= CheckSupportRule(TypeAnyOf(output, supportedTypes), reasonIfUnsupported,
+                                  "Reference poolind2d: output is not a supported type.");
+
+    supported &= CheckSupportRule(TypesAreEqual(input, output), reasonIfUnsupported,
+                                  "Reference poolind2d: input and output types are mismatched.");
+
+    return supported;
 }
 
 bool RefLayerSupport::IsQuantizeSupported(const TensorInfo& input,
-- 
cgit v1.2.1