commit 186c21c9a598cbe2e81ad53e5b5fd96d75f981f5
author Francis Murtagh <francis.murtagh@arm.com> Thu Aug 20 15:38:29 2020 +0100
committer Teresa Charlin <teresa.charlinreyes@arm.com> Fri Aug 21 15:56:13 2020 +0100
tree 45cb5522dcf36256f66bbe25eb8b24460cbdb37a
parent c84e45d933a9b45810a3bb88f6873f4eddca0975

Bugfix: Allow permutation of QuantizationDim

Signed-off-by: Francis Murtagh <francis.murtagh@arm.com>
Change-Id: Ib98ec471e6fdd47600b7c62d0b4d19dd36e20cbd

include/armnnUtils/Permute.hpp

diff --git a/include/armnnUtils/Permute.hpp b/include/armnnUtils/Permute.hpp
index 4e9bfc0..d719f4a 100644
--- a/include/armnnUtils/Permute.hpp
+++ b/include/armnnUtils/Permute.hpp
@@ -11,9 +11,12 @@
 namespace armnnUtils
 {
 
-armnn::TensorShape Permuted(const armnn::TensorShape& srcShape, const armnn::PermutationVector& mappings);
+armnn::TensorShape Permuted(const armnn::TensorShape& srcShape,
+                            const armnn::PermutationVector& mappings);
 
-armnn::TensorInfo Permuted(const armnn::TensorInfo& info, const armnn::PermutationVector& mappings);
+armnn::TensorInfo Permuted(const armnn::TensorInfo& info,
+                           const armnn::PermutationVector& mappings,
+                           bool perChannelPermute = false);
 
 void Permute(const armnn::TensorShape& dstShape, const armnn::PermutationVector& mappings,
              const void* src, void* dst, size_t dataTypeSize);

src/armnn/test/UtilsTests.cpp

diff --git a/src/armnn/test/UtilsTests.cpp b/src/armnn/test/UtilsTests.cpp
index 0bae756..fb078de 100644
--- a/src/armnn/test/UtilsTests.cpp
+++ b/src/armnn/test/UtilsTests.cpp
@@ -9,6 +9,7 @@
 #include <armnn/Types.hpp>
 #include <armnn/TypesUtils.hpp>
 #include <armnn/Descriptors.hpp>
+#include <armnnUtils/Permute.hpp>
 #include <GraphTopologicalSort.hpp>
 #include <Graph.hpp>
 #include <ResolveType.hpp>
@@ -245,4 +246,24 @@
     BOOST_TEST(!sortCompleted);
 }
 
+BOOST_AUTO_TEST_CASE(PermuteQuantizationDim)
+{
+    std::vector<float> scales;
+
+    // Set QuantizationDim to be index 1
+    const armnn::TensorInfo info({ 1, 2, 3, 4 }, armnn::DataType::Float32, scales, 1U);
+    BOOST_CHECK(info.GetQuantizationDim().value() == 1U);
+
+    // Permute so that index 1 moves to final index i.e. index 3
+    armnn::PermutationVector mappings({ 0, 3, 2, 1 });
+    auto permutedPerChannel = armnnUtils::Permuted(info, mappings, true);
+    auto permuted = armnnUtils::Permuted(info, mappings);
+
+    // Check that QuantizationDim is in index 3
+    BOOST_CHECK(permutedPerChannel.GetQuantizationDim().value() == 3U);
+
+    // Check previous implementation unchanged
+    BOOST_CHECK(permuted.GetQuantizationDim().value() == 1U);
+}
+
 BOOST_AUTO_TEST_SUITE_END()

src/armnnUtils/Permute.cpp

diff --git a/src/armnnUtils/Permute.cpp b/src/armnnUtils/Permute.cpp
index 486aac0..3770463 100644
--- a/src/armnnUtils/Permute.cpp
+++ b/src/armnnUtils/Permute.cpp
@@ -95,7 +95,8 @@
 namespace armnnUtils
 {
 
-armnn::TensorShape Permuted(const armnn::TensorShape& srcShape, const armnn::PermutationVector& mappings)
+armnn::TensorShape Permuted(const armnn::TensorShape& srcShape,
+                            const armnn::PermutationVector& mappings)
 {
     assert(srcShape.GetNumDimensions() == mappings.GetSize());
 
@@ -111,10 +112,19 @@
     return permutedShape;
 }
 
-armnn::TensorInfo Permuted(const armnn::TensorInfo& info, const armnn::PermutationVector& mappings)
+armnn::TensorInfo Permuted(const armnn::TensorInfo& info,
+                           const armnn::PermutationVector& mappings,
+                           bool perChannelPermute)
 {
     armnn::TensorInfo outInfo(info);
     outInfo.SetShape(Permuted(info.GetShape(), mappings));
+
+    // If TensorInfo has Per-Axis Quantization then permute QuantizationDim to mapping
+    if (info.HasPerAxisQuantization() && perChannelPermute)
+    {
+        outInfo.SetQuantizationDim(mappings[info.GetQuantizationDim().value()]);
+    }
+
     return outInfo;
 }