IVGCVSW-3831 Add support of per-axis quantization to TensorInfo

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

3 files changed, 199 insertions, 40 deletions

diff --git a/include/armnn/Tensor.hpp b/include/armnn/Tensor.hpp
index b3a46290ae..57a243800e 100644
--- a/include/armnn/Tensor.hpp
+++ b/include/armnn/Tensor.hpp
@@ -7,6 +7,7 @@
 #include "TensorFwd.hpp"
 
 #include "Exceptions.hpp"
+#include "Optional.hpp"
 #include "Types.hpp"
 
 #include <array>
@@ -55,10 +56,27 @@ public:
     /// Empty (invalid) constructor.
     TensorInfo();
 
-    TensorInfo(const TensorShape& shape, DataType dataType,
-        float quantizationScale = 0.0f, int32_t quantizationOffset = 0);
-    TensorInfo(unsigned int numDimensions, const unsigned int* dimensionSizes, DataType dataType,
-        float quantizationScale = 0.0f, int32_t quantizationOffset = 0);
+    TensorInfo(const TensorShape& shape,
+               DataType dataType,
+               float quantizationScale = 0.0f,
+               int32_t quantizationOffset = 0);
+
+    TensorInfo(unsigned int numDimensions,
+               const unsigned int* dimensionSizes,
+               DataType dataType,
+               float quantizationScale = 0.0f,
+               int32_t quantizationOffset = 0);
+
+    TensorInfo(const TensorShape& shape,
+               DataType dataType,
+               const std::vector<float>& quantizationScales,
+               unsigned int quantizationDim);
+
+    TensorInfo(unsigned int numDimensions,
+               const unsigned int* dimensionSizes,
+               DataType dataType,
+               const std::vector<float>& quantizationScales,
+               unsigned int quantizationDim);
 
     TensorInfo(const TensorInfo& other);
 
@@ -67,22 +85,31 @@ public:
     bool operator==(const TensorInfo& other) const;
     bool operator!=(const TensorInfo& other) const;
 
-    const TensorShape& GetShape() const             { return m_Shape; }
-    TensorShape& GetShape()                         { return m_Shape; }
-    void SetShape(const TensorShape& newShape)      { m_Shape = newShape; }
+    const TensorShape& GetShape() const              { return m_Shape; }
+    TensorShape& GetShape()                          { return m_Shape; }
+    void SetShape(const TensorShape& newShape)       { m_Shape = newShape; }
+
+    unsigned int GetNumDimensions() const            { return m_Shape.GetNumDimensions(); }
+    unsigned int GetNumElements() const              { return m_Shape.GetNumElements(); }
+
+    DataType GetDataType() const                     { return m_DataType; }
+    void SetDataType(DataType type)                  { m_DataType = type; }
+
+    bool HasMultipleQuantizationScales() const       { return m_Quantization.m_Scales.size() > 1; }
+
+    std::vector<float> GetQuantizationScales() const;
+    void SetQuantizationScales(const std::vector<float>& scales);
 
-    unsigned int GetNumDimensions() const { return m_Shape.GetNumDimensions(); }
-    unsigned int GetNumElements() const { return m_Shape.GetNumElements(); }
+    float GetQuantizationScale() const;
+    void SetQuantizationScale(float scale);
 
-    DataType GetDataType() const                    { return m_DataType; }
-    void SetDataType(DataType type)                 { m_DataType = type; }
+    int32_t GetQuantizationOffset() const;
+    void SetQuantizationOffset(int32_t offset);
 
-    float GetQuantizationScale() const              { return m_Quantization.m_Scale; }
-    int32_t GetQuantizationOffset() const           { return m_Quantization.m_Offset; }
-    void SetQuantizationScale(float scale)          { m_Quantization.m_Scale = scale; }
-    void SetQuantizationOffset(int32_t offset)      { m_Quantization.m_Offset = offset; }
-    bool IsQuantized() const                        { return m_DataType == DataType::QuantisedAsymm8 ||
-                                                             m_DataType == DataType::QuantisedSymm16; }
+    Optional<unsigned int> GetQuantizationDim() const;
+    void SetQuantizationDim(const Optional<unsigned int>& quantizationDim);
+
+    bool IsQuantized() const;
 
     /// Check that the types are the same and, if quantize, that the quantization parameters are the same.
     bool IsTypeSpaceMatch(const TensorInfo& other) const;
@@ -91,14 +118,26 @@ public:
 
 private:
     TensorShape m_Shape;
-    DataType m_DataType;
-    /// Scale and offset values are used for quantization.
+    DataType    m_DataType;
+
+    /// Vectors of scale and offset are used for per-axis quantization.
     struct Quantization
     {
-        Quantization() : m_Scale(0.f), m_Offset(0) {}
-        bool operator==(const Quantization& o) const {return ((m_Scale == o.m_Scale) && (m_Offset == o.m_Offset));}
-        float m_Scale;
-        int32_t m_Offset;
+        Quantization()
+            : m_Scales{}
+            , m_Offset(EmptyOptional())
+            , m_QuantizationDim(EmptyOptional()) {}
+
+        bool operator==(const Quantization& other) const
+        {
+            return ((m_Scales == other.m_Scales) && (m_Offset == other.m_Offset) &&
+                (m_QuantizationDim == other.m_QuantizationDim));
+        }
+
+        std::vector<float>     m_Scales;
+        Optional<int32_t>      m_Offset;
+        Optional<unsigned int> m_QuantizationDim;
+
     } m_Quantization;
 };
 
@@ -151,7 +190,7 @@ class Tensor : public BaseTensor<void*>
 {
 public:
     /// Brings in the constructors and assignment operator.
-    using BaseTensor<void*>::BaseTensor; 
+    using BaseTensor<void*>::BaseTensor;
 };
 
 /// A tensor defined by a TensorInfo (shape and data type) and an immutable backing store.
@@ -159,7 +198,7 @@ class ConstTensor : public BaseTensor<const void*>
 {
 public:
     /// Brings in the constructors and assignment operator.
-    using BaseTensor<const void*>::BaseTensor; 
+    using BaseTensor<const void*>::BaseTensor;
     ConstTensor() : BaseTensor<const void*>() {} // This needs to be redefined explicitly??
 
     /// Can be implicitly constructed from non-const Tensor.
diff --git a/src/armnn/Tensor.cpp b/src/armnn/Tensor.cpp
index 614abc77f5..f4b8b509b6 100644
--- a/src/armnn/Tensor.cpp
+++ b/src/armnn/Tensor.cpp
@@ -2,6 +2,7 @@
 // Copyright © 2017 Arm Ltd. All rights reserved.
 // SPDX-License-Identifier: MIT
 //
+
 #include "armnn/Tensor.hpp"
 #include "armnn/Utils.hpp"
 #include "armnn/Exceptions.hpp"
@@ -138,30 +139,57 @@ TensorInfo::TensorInfo()
 {
 }
 
-TensorInfo::TensorInfo(const TensorShape& shape, DataType dataType,
-    float quantizationScale, int32_t quantizationOffset)
- : m_Shape(shape)
- , m_DataType(dataType)
+TensorInfo::TensorInfo(const TensorShape& shape,
+                       DataType dataType,
+                       float quantizationScale,
+                       int32_t quantizationOffset)
+    : m_Shape(shape)
+    , m_DataType(dataType)
+{
+    SetQuantizationScale(quantizationScale);
+    SetQuantizationOffset(quantizationOffset);
+}
+
+TensorInfo::TensorInfo(unsigned int numDimensions,
+                       const unsigned int* dimensionSizes,
+                       DataType dataType,
+                       float quantizationScale,
+                       int32_t quantizationOffset)
+    : m_Shape(numDimensions, dimensionSizes)
+    , m_DataType(dataType)
 {
-    m_Quantization.m_Scale = quantizationScale;
-    m_Quantization.m_Offset = quantizationOffset;
+    SetQuantizationScale(quantizationScale);
+    SetQuantizationOffset(quantizationOffset);
 }
 
-TensorInfo::TensorInfo(unsigned int numDimensions, const unsigned int* dimensionSizes, DataType dataType,
-    float quantizationScale, int32_t quantizationOffset)
- : m_Shape(numDimensions, dimensionSizes)
- , m_DataType(dataType)
+TensorInfo::TensorInfo(const TensorShape& shape,
+                       DataType dataType,
+                       const std::vector<float>& quantizationScales,
+                       unsigned int quantizationDim)
+    : m_Shape(shape)
+    , m_DataType(dataType)
 {
-    m_Quantization.m_Scale = quantizationScale;
-    m_Quantization.m_Offset = quantizationOffset;
+    SetQuantizationScales(quantizationScales);
+    SetQuantizationDim(MakeOptional<unsigned int>(quantizationDim));
+}
+
+TensorInfo::TensorInfo(unsigned int numDimensions,
+                       const unsigned int* dimensionSizes,
+                       DataType dataType,
+                       const std::vector<float>& quantizationScales,
+                       unsigned int quantizationDim)
+    : m_Shape(numDimensions, dimensionSizes)
+    , m_DataType(dataType)
+{
+    SetQuantizationScales(quantizationScales);
+    SetQuantizationDim(MakeOptional<unsigned int>(quantizationDim));
 }
 
 TensorInfo::TensorInfo(const TensorInfo& other)
 : m_Shape(other.m_Shape)
 , m_DataType(other.m_DataType)
 , m_Quantization(other.m_Quantization)
-{
-}
+{}
 
 TensorInfo& TensorInfo::operator=(const TensorInfo& other)
 {
@@ -194,7 +222,7 @@ bool TensorInfo::IsTypeSpaceMatch(const TensorInfo& other) const
 
     match &= m_DataType == other.m_DataType;
 
-    if (IsQuantized())
+    if (IsQuantized() && !HasMultipleQuantizationScales())
     {
         match &= GetQuantizationScale() == other.GetQuantizationScale() &&
                  GetQuantizationOffset() == other.GetQuantizationOffset();
@@ -202,6 +230,64 @@ bool TensorInfo::IsTypeSpaceMatch(const TensorInfo& other) const
     return match;
 }
 
+std::vector<float> TensorInfo::GetQuantizationScales() const
+{
+    return m_Quantization.m_Scales;
+}
+
+void TensorInfo::SetQuantizationScales(const std::vector<float>& scales)
+{
+    m_Quantization.m_Scales = scales;
+}
+
+float TensorInfo::GetQuantizationScale() const
+{
+    if (m_Quantization.m_Scales.empty())
+    {
+        // NOTE: old default for backward compatibility
+        return 1.0f;
+    }
+
+    BOOST_ASSERT(!HasMultipleQuantizationScales());
+    return m_Quantization.m_Scales[0];
+}
+
+void TensorInfo::SetQuantizationScale(float scale)
+{
+    m_Quantization.m_Scales = { scale };
+}
+
+int32_t TensorInfo::GetQuantizationOffset() const
+{
+    if (!m_Quantization.m_Offset.has_value())
+    {
+        // NOTE: old default for backward compatibility
+        return 0;
+    }
+
+    return m_Quantization.m_Offset.value();
+}
+
+void TensorInfo::SetQuantizationOffset(int32_t offset)
+{
+    m_Quantization.m_Offset = MakeOptional<int32_t>(offset);
+}
+
+Optional<unsigned int> TensorInfo::GetQuantizationDim() const
+{
+    return m_Quantization.m_QuantizationDim;
+}
+
+void TensorInfo::SetQuantizationDim(const Optional<unsigned int>& quantizationDim)
+{
+    m_Quantization.m_QuantizationDim = quantizationDim;
+}
+
+bool TensorInfo::IsQuantized() const
+{
+    return m_DataType == DataType::QuantisedAsymm8 || m_DataType == DataType::QuantisedSymm16;
+}
+
 // ---
 // --- BaseTensor
 // ---
diff --git a/src/armnn/test/TensorTest.cpp b/src/armnn/test/TensorTest.cpp
index a0a6c7e91f..154a0bca04 100644
--- a/src/armnn/test/TensorTest.cpp
+++ b/src/armnn/test/TensorTest.cpp
@@ -143,4 +143,38 @@ BOOST_AUTO_TEST_CASE(TensorShapeOperatorBrackets)
     BOOST_TEST(shape[2] == 20);
 }
 
+BOOST_AUTO_TEST_CASE(TensorInfoPerAxisQuantization)
+{
+    // Old constructor
+    TensorInfo tensorInfo0({ 1, 1 }, DataType::Float32, 2.0f, 1);
+    BOOST_CHECK(!tensorInfo0.HasMultipleQuantizationScales());
+    BOOST_CHECK(tensorInfo0.GetQuantizationScale() == 2.0f);
+    BOOST_CHECK(tensorInfo0.GetQuantizationOffset() == 1);
+    BOOST_CHECK(tensorInfo0.GetQuantizationScales()[0] == 2.0f);
+    BOOST_CHECK(!tensorInfo0.GetQuantizationDim().has_value());
+
+    // Set per-axis quantization scales
+    std::vector<float> perAxisScales{ 3.0f, 4.0f };
+    tensorInfo0.SetQuantizationScales(perAxisScales);
+    BOOST_CHECK(tensorInfo0.HasMultipleQuantizationScales());
+    BOOST_CHECK(tensorInfo0.GetQuantizationScales() == perAxisScales);
+
+    // Set per-tensor quantization scale
+    tensorInfo0.SetQuantizationScale(5.0f);
+    BOOST_CHECK(!tensorInfo0.HasMultipleQuantizationScales());
+    BOOST_CHECK(tensorInfo0.GetQuantizationScales()[0] == 5.0f);
+
+    // Set quantization offset
+    tensorInfo0.SetQuantizationDim(Optional<unsigned int>(1));
+    BOOST_CHECK(tensorInfo0.GetQuantizationDim().value() == 1);
+
+    // New constructor
+    perAxisScales = { 6.0f, 7.0f };
+    TensorInfo tensorInfo1({ 1, 1 }, DataType::Float32, perAxisScales, 1);
+    BOOST_CHECK(tensorInfo1.HasMultipleQuantizationScales());
+    BOOST_CHECK(tensorInfo1.GetQuantizationOffset() == 0);
+    BOOST_CHECK(tensorInfo1.GetQuantizationScales() == perAxisScales);
+    BOOST_CHECK(tensorInfo1.GetQuantizationDim().value() == 1);
+}
+
 BOOST_AUTO_TEST_SUITE_END()