IVGCVSW-3842 Add per-axis quantization support to ValidateBiasTensorQuantization()

* Altered the function to validate all the scale values in the bias
  tensor info when per-axis quantization is enabled
* Added unit test to check if ValidateBiasTensorQuantization() works
  as intended

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

2 files changed, 105 insertions, 8 deletions

diff --git a/src/backends/backendsCommon/WorkloadData.cpp b/src/backends/backendsCommon/WorkloadData.cpp
index cfb38b4820..e3d78fea50 100644
--- a/src/backends/backendsCommon/WorkloadData.cpp
+++ b/src/backends/backendsCommon/WorkloadData.cpp
@@ -192,20 +192,52 @@ void ValidateBiasTensorQuantization(const TensorInfo& biasTensor,
                                     const TensorInfo& weightsTensorInfo,
                                     const std::string& descName)
 {
+    // Helper lambda function to validate a single bias quantization scale value
+    auto VerifyBiasQuantizationScale = [&descName](float biasScale, float expectedScale) -> void
+    {
+        constexpr float tolerance = 0.00000001f;
+        if (std::abs(biasScale - expectedScale) > tolerance)
+        {
+            // Print the float values with extra precision to see very small differences
+            std::stringstream msg;
+            msg << std::setprecision(10) << descName << ": Expected " << expectedScale <<
+                " quantization scale for bias tensor (the product of the input and weight scales), but got " <<
+                biasScale;
+            throw InvalidArgumentException(msg.str(), CHECK_LOCATION());
+        }
+    };
+
     if (biasTensor.GetQuantizationOffset() != 0)
     {
         throw InvalidArgumentException(descName + ": Expected zero quantization offset for bias tensor but got " +
             to_string(biasTensor.GetQuantizationOffset()));
     }
-    const float expectedScale = inputTensorInfo.GetQuantizationScale() * weightsTensorInfo.GetQuantizationScale();
-    if (std::abs(biasTensor.GetQuantizationScale() - expectedScale) > 0.00000001f)
+
+    if (biasTensor.HasMultipleQuantizationScales())
+    {
+        // Validate per-axis quantization scales
+        const std::vector<float>& weightScales = weightsTensorInfo.GetQuantizationScales();
+        const std::vector<float>& biasScales   = biasTensor.GetQuantizationScales();
+
+        if (weightScales.size() != biasScales.size())
+        {
+            std::stringstream msg;
+            msg << descName << ": Expected matchhing number of per-axis quantization scales, but got different "
+                << "values: weights=" << weightScales.size()  << ", biases=" << biasScales.size();
+            throw InvalidArgumentException(msg.str(), CHECK_LOCATION());
+        }
+
+        for (size_t i = 0ul; i < biasScales.size(); ++i)
+        {
+            const float expectedScale = inputTensorInfo.GetQuantizationScale() * weightScales[i];
+            VerifyBiasQuantizationScale(biasScales[i], expectedScale);
+        }
+    }
+    else
     {
-        // Print the float values with extra precision to see very small differences
-        std::stringstream msg;
-        msg << std::setprecision(10) << descName << ": Expected " << expectedScale <<
-            " quantization scale for bias tensor (the product of the input and weight scales), but got " <<
-            biasTensor.GetQuantizationScale();
-        throw InvalidArgumentException(msg.str());
+        // Validate per-tensor quantization scale
+        const float expectedScale = inputTensorInfo.GetQuantizationScale() * weightsTensorInfo.GetQuantizationScale();
+        VerifyBiasQuantizationScale(biasTensor.GetQuantizationScale(), expectedScale);
     }
 }
 
diff --git a/src/backends/backendsCommon/test/WorkloadDataValidation.cpp b/src/backends/backendsCommon/test/WorkloadDataValidation.cpp
index c6960986b3..9773914220 100644
--- a/src/backends/backendsCommon/test/WorkloadDataValidation.cpp
+++ b/src/backends/backendsCommon/test/WorkloadDataValidation.cpp
@@ -588,4 +588,69 @@ BOOST_AUTO_TEST_CASE(LstmQueueDescriptor_Validate)
     BOOST_CHECK_NO_THROW(data.Validate(info));
 }
 
+BOOST_AUTO_TEST_CASE(BiasPerAxisQuantization_Validate)
+{
+    constexpr unsigned int nInput  = 1u;
+    constexpr unsigned int cInput  = 3u;
+    constexpr unsigned int hInput  = 3u;
+    constexpr unsigned int wInput  = 3u;
+
+    constexpr unsigned int nOutput = nInput;
+    constexpr unsigned int cOutput = cInput;
+    constexpr unsigned int hOutput = 1u;
+    constexpr unsigned int wOutput = 1u;
+
+    const TensorShape inputShape { nInput,  cInput,  hInput,  wInput  };
+    const TensorShape outputShape{ nOutput, cOutput, hOutput, wOutput };
+    const TensorShape weightShape{ cOutput, cInput,  hInput,  wInput  };
+    const TensorShape biasShape  { cOutput                            };
+
+    constexpr DataType dataType = DataType::QuantisedAsymm8;
+    constexpr DataType biasType = DataType::Signed32;
+
+    constexpr float perTensorScale = 1.5f;
+    const TensorInfo inputInfo (inputShape,  dataType, perTensorScale);
+    const TensorInfo outputInfo(outputShape, dataType, perTensorScale);
+
+    const std::vector<float> weightPerAxisScales = { 2.50f, 3.50f };
+    const TensorInfo weightInfo(weightShape, dataType, weightPerAxisScales, 0);
+
+    Convolution2dQueueDescriptor queueDescriptor;
+    queueDescriptor.m_Parameters.m_BiasEnabled = true;
+
+    WorkloadInfo workloadInfo;
+    AddInputToWorkload(queueDescriptor, workloadInfo, inputInfo, nullptr);
+    AddOutputToWorkload(queueDescriptor, workloadInfo, outputInfo, nullptr);
+
+    ScopedCpuTensorHandle weightTensor(weightInfo);
+    queueDescriptor.m_Weight = &weightTensor;
+
+    // Test 1: correct per-axis quantization values
+    const std::vector<float> biasPerAxisScales1  = { 3.75f, 5.25f };
+    const TensorInfo biasInfo1(biasShape, biasType, biasPerAxisScales1, 0);
+
+    ScopedCpuTensorHandle biasHandle1(biasInfo1);
+    queueDescriptor.m_Bias = &biasHandle1;
+
+    BOOST_CHECK_NO_THROW(queueDescriptor.Validate(workloadInfo));
+
+    // Test 2: wrong per-axis quantization values
+    const std::vector<float> biasPerAxisScales2 = { 4.00f, 5.00f };
+    const TensorInfo biasInfo2(biasShape, biasType, biasPerAxisScales2, 0);
+
+    ScopedCpuTensorHandle biasHandle2(biasInfo2);
+    queueDescriptor.m_Bias = &biasHandle2;
+
+    BOOST_CHECK_THROW(queueDescriptor.Validate(workloadInfo), InvalidArgumentException);
+
+    // Test 3: mismatched number of quantization scales
+    const std::vector<float> biasPerAxisScales3 = { 3.75f, 5.25f, 5.25f };
+    const TensorInfo biasInfo3(biasShape, biasType, biasPerAxisScales3, 0);
+
+    ScopedCpuTensorHandle biasHandle3(biasInfo3);
+    queueDescriptor.m_Bias = &biasHandle3;
+
+    BOOST_CHECK_THROW(queueDescriptor.Validate(workloadInfo), InvalidArgumentException);
+}
+
 BOOST_AUTO_TEST_SUITE_END()