IVGCVSW-3839 Add support of per-axis quantization to reference TransposeConvolution2d

Signed-off-by: Aron Virginas-Tar <Aron.Virginas-Tar@arm.com>
Change-Id: Ie0dc1204eee925adfb1e59aba3f1137178302184

3 files changed, 31 insertions, 7 deletions

diff --git a/src/backends/reference/RefLayerSupport.cpp b/src/backends/reference/RefLayerSupport.cpp
index 4252fecf35..e98af7097b 100644
--- a/src/backends/reference/RefLayerSupport.cpp
+++ b/src/backends/reference/RefLayerSupport.cpp
@@ -1802,14 +1802,31 @@ bool RefLayerSupport::IsTransposeConvolution2dSupported(const TensorInfo& input,
     supported &= CheckSupportRule(TypeAnyOf(output, supportedTypes), reasonIfUnsupported,
                                   "Reference TransposeConvolution2d: output is not a supported type.");
 
-    supported &= CheckSupportRule(TypeAnyOf(weights, supportedTypes), reasonIfUnsupported,
-                                  "Reference TransposeConvolution2d: weights is not a supported type.");
-
     supported &= CheckSupportRule(TypesAreEqual(input, output), reasonIfUnsupported,
                                   "Reference TransposeConvolution2d: input and output types mismatched.");
 
-    supported &= CheckSupportRule(TypesAreEqual(input, weights), reasonIfUnsupported,
-                                  "Reference TransposeConvolution2d: input and weights types mismatched.");
+
+    const DataType inputType = input.GetDataType();
+    if (inputType == DataType::QuantisedAsymm8)
+    {
+        std::array<DataType, 2> supportedWeightTypes =
+        {
+            DataType::QuantisedAsymm8,
+            DataType::QuantizedSymm8PerAxis
+        };
+
+        supported &= CheckSupportRule(TypeAnyOf(weights, supportedWeightTypes), reasonIfUnsupported,
+                                      "Reference TransposeConvolution2d: weights type not supported for "
+                                      "quantized input.");
+    }
+    else
+    {
+        supported &= CheckSupportRule(TypeAnyOf(weights, supportedTypes), reasonIfUnsupported,
+                                    "Reference TransposeConvolution2d: weights is not a supported type.");
+
+        supported &= CheckSupportRule(TypesAreEqual(input, weights), reasonIfUnsupported,
+                                    "Reference TransposeConvolution2d: input and weights types mismatched.");
+    }
 
     if (biases.has_value())
     {
diff --git a/src/backends/reference/test/RefLayerTests.cpp b/src/backends/reference/test/RefLayerTests.cpp
index c407828f06..cd5c9273f7 100644
--- a/src/backends/reference/test/RefLayerTests.cpp
+++ b/src/backends/reference/test/RefLayerTests.cpp
@@ -1609,6 +1609,13 @@ ARMNN_AUTO_TEST_CASE(MultiChannelTransposeConvolution2dInt16Nhwc,
                      MultiChannelTransposeConvolution2dTest<DataType::QuantisedSymm16, DataType::Signed32>,
                      DataLayout::NCHW)
 
+ARMNN_AUTO_TEST_CASE(TransposeConvolution2dPerAxisQuantTestNchw,
+                     TransposeConvolution2dPerAxisQuantTest,
+                     DataLayout::NCHW);
+ARMNN_AUTO_TEST_CASE(TransposeConvolution2dPerAxisQuantTestNhwc,
+                     TransposeConvolution2dPerAxisQuantTest,
+                     DataLayout::NHWC);
+
 // Stack
 ARMNN_AUTO_TEST_CASE(Stack0Axis,           StackAxis0Float32Test)
 ARMNN_AUTO_TEST_CASE(StackOutput4DAxis1,   StackOutput4DAxis1Float32Test)
diff --git a/src/backends/reference/workloads/TransposeConvolution2d.cpp b/src/backends/reference/workloads/TransposeConvolution2d.cpp
index 1b89ef9262..5662c58809 100644
--- a/src/backends/reference/workloads/TransposeConvolution2d.cpp
+++ b/src/backends/reference/workloads/TransposeConvolution2d.cpp
@@ -79,7 +79,7 @@ void TransposeConvolution2dImpl(const TransposeConvolution2dDescriptor& descript
 
                                     const unsigned int weightsIndex =
                                         dataLayoutIndexed.GetIndex(weightsShape, dOutput, dInput, yWeights, xWeights);
-                                    weightsDecoder[weightsIndex];
+                                    weightsDecoder.SetIndex(weightsIndex, dOutput);
 
                                     const unsigned int outputIndex =
                                         dataLayoutIndexed.GetIndex(outputShape, batch, dOutput, yOutput, xOutput);
@@ -107,7 +107,7 @@ void TransposeConvolution2dImpl(const TransposeConvolution2dDescriptor& descript
         {
             for (unsigned int dOutput = 0u; dOutput < outputDepth; ++dOutput)
             {
-                rBiasesDecoder[dOutput];
+                rBiasesDecoder.SetIndex(dOutput, dOutput);
                 for (unsigned int yOutput = 0u; yOutput < outputHeight; ++yOutput)
                 {
                     for (unsigned int xOutput = 0u; xOutput < outputWidth; ++xOutput)