IVGCVSW-3705 Add Channel Shuffle Front end and Ref Implementation

* Add front end
* Add reference workload
* Add unit tests
* Add Serializer and Deserializer
* Update ArmNN Versioning

Signed-off-by: Simon Obute <simon.obute@arm.com>
Change-Id: I9ac1f953af3974382eac8e8d62d794d2344e8f47

10 files changed, 167 insertions, 0 deletions

diff --git a/src/backends/reference/RefLayerSupport.cpp b/src/backends/reference/RefLayerSupport.cpp
index 5eba3e5919..aaf9aa0e7c 100644
--- a/src/backends/reference/RefLayerSupport.cpp
+++ b/src/backends/reference/RefLayerSupport.cpp
@@ -333,6 +333,39 @@ bool RefLayerSupport::IsCastSupported(const TensorInfo& input,
     return supported;
 }
 
+bool RefLayerSupport::IsChannelShuffleSupported(const TensorInfo& input,
+                                                const TensorInfo& output,
+                                                const ChannelShuffleDescriptor& descriptor,
+                                                Optional<std::string&> reasonIfUnsupported) const
+{
+    IgnoreUnused(descriptor);
+    bool supported = true;
+
+    // Define supported output and inputs types.
+    std::array<DataType, 7> supportedTypes =
+    {
+        DataType::BFloat16,
+        DataType::Float32,
+        DataType::Float16,
+        DataType::QAsymmS8,
+        DataType::QAsymmU8,
+        DataType::QSymmS8,
+        DataType::QSymmS16
+    };
+
+    supported &= CheckSupportRule(TypeAnyOf(input, supportedTypes), reasonIfUnsupported,
+                                  "Reference ChannelShuffle: input is not a supported type.");
+
+    supported &= CheckSupportRule(TypeAnyOf(output, supportedTypes), reasonIfUnsupported,
+                                  "Reference ChannelShuffle: output is not a supported type.");
+
+    supported &= CheckSupportRule(TypesAreEqual(input, output), reasonIfUnsupported,
+                                  "Reference ChannelShuffle: input and output types are mismatched.");
+
+    return supported;
+}
+
+
 bool RefLayerSupport::IsComparisonSupported(const TensorInfo& input0,
                                             const TensorInfo& input1,
                                             const TensorInfo& output,
diff --git a/src/backends/reference/RefLayerSupport.hpp b/src/backends/reference/RefLayerSupport.hpp
index a1b4dc7f47..2693dc1779 100644
--- a/src/backends/reference/RefLayerSupport.hpp
+++ b/src/backends/reference/RefLayerSupport.hpp
@@ -50,6 +50,11 @@ public:
                          const TensorInfo& output,
                          Optional<std::string&> reasonIfUnsupported = EmptyOptional()) const override;
 
+    bool IsChannelShuffleSupported(const TensorInfo& input,
+                                   const TensorInfo& output,
+                                   const ChannelShuffleDescriptor& descriptor,
+                                   Optional<std::string&> reasonIfUnsupported = EmptyOptional()) const override;
+
     bool IsComparisonSupported(const TensorInfo& input0,
                                const TensorInfo& input1,
                                const TensorInfo& output,
diff --git a/src/backends/reference/RefWorkloadFactory.cpp b/src/backends/reference/RefWorkloadFactory.cpp
index 16cf17cc79..681b73a748 100644
--- a/src/backends/reference/RefWorkloadFactory.cpp
+++ b/src/backends/reference/RefWorkloadFactory.cpp
@@ -183,6 +183,12 @@ std::unique_ptr<IWorkload> RefWorkloadFactory::CreateCast(const CastQueueDescrip
     return std::make_unique<RefCastWorkload>(descriptor, info);
 }
 
+std::unique_ptr<IWorkload> RefWorkloadFactory::CreateChannelShuffle(const ChannelShuffleQueueDescriptor &descriptor,
+                                                                    const WorkloadInfo &info) const
+{
+    return std::make_unique<RefChannelShuffleWorkload>(descriptor,info);
+}
+
 std::unique_ptr<IWorkload> RefWorkloadFactory::CreateComparison(const ComparisonQueueDescriptor& descriptor,
                                                                 const WorkloadInfo& info) const
 {
diff --git a/src/backends/reference/RefWorkloadFactory.hpp b/src/backends/reference/RefWorkloadFactory.hpp
index 113aca70ef..fe3eb54795 100644
--- a/src/backends/reference/RefWorkloadFactory.hpp
+++ b/src/backends/reference/RefWorkloadFactory.hpp
@@ -88,6 +88,9 @@ public:
     std::unique_ptr<IWorkload> CreateCast(const CastQueueDescriptor& descriptor,
                                           const WorkloadInfo& info) const override;
 
+    std::unique_ptr<IWorkload> CreateChannelShuffle(const ChannelShuffleQueueDescriptor& descriptor,
+                                                    const WorkloadInfo& info) const override;
+
     std::unique_ptr<IWorkload> CreateComparison(const ComparisonQueueDescriptor& descriptor,
                                                 const WorkloadInfo& info) const override;
 
diff --git a/src/backends/reference/backend.mk b/src/backends/reference/backend.mk
index 17ddbe0df1..2dc2bc4919 100644
--- a/src/backends/reference/backend.mk
+++ b/src/backends/reference/backend.mk
@@ -49,6 +49,7 @@ BACKEND_SOURCES := \
         workloads/RefBatchNormalizationWorkload.cpp \
         workloads/RefBatchToSpaceNdWorkload.cpp \
         workloads/RefCastWorkload.cpp \
+        workloads/RefChannelShuffleWorkload.cpp \
         workloads/RefComparisonWorkload.cpp \
         workloads/RefConcatWorkload.cpp \
         workloads/RefConstantWorkload.cpp \
diff --git a/src/backends/reference/test/RefLayerTests.cpp b/src/backends/reference/test/RefLayerTests.cpp
index 18490e29c7..b0855151de 100644
--- a/src/backends/reference/test/RefLayerTests.cpp
+++ b/src/backends/reference/test/RefLayerTests.cpp
@@ -1398,6 +1398,13 @@ ARMNN_AUTO_TEST_CASE_WITH_THF(SimpleReshapeQuantisedSymm16, SimpleReshapeTest<Da
 ARMNN_AUTO_TEST_CASE_WITH_THF(Reshape5d, Reshape5dTest<DataType::Float32>)
 ARMNN_AUTO_TEST_CASE_WITH_THF(ReshapeBoolean, ReshapeBooleanTest)
 
+// ChannelShuffle
+ARMNN_AUTO_TEST_CASE_WITH_THF(SimpleChannelShuffleFloat32, SimpleChannelShuffleTest<DataType::Float32>)
+ARMNN_AUTO_TEST_CASE_WITH_THF(SimpleChannelShuffleQAsymmU8, SimpleChannelShuffleTest<DataType::QAsymmU8>)
+ARMNN_AUTO_TEST_CASE_WITH_THF(ChannelShuffle2DFloat32, ChannelShuffle2DTest<DataType::Float32>)
+ARMNN_AUTO_TEST_CASE_WITH_THF(ChannelShuffle2DQAsymmU8, ChannelShuffle2DTest<DataType::QAsymmU8>)
+ARMNN_AUTO_TEST_CASE_WITH_THF(ChannelShuffle4DFloat32, ChannelShuffle4DTest<DataType::Float32>)
+ARMNN_AUTO_TEST_CASE_WITH_THF(ChannelShuffle4DQAsymmU8, ChannelShuffle4DTest<DataType::QAsymmU8>)
 
 // Rsqrt
 ARMNN_AUTO_TEST_CASE_WITH_THF(Rsqrt2d, Rsqrt2dTest<DataType::Float32>)
diff --git a/src/backends/reference/workloads/CMakeLists.txt b/src/backends/reference/workloads/CMakeLists.txt
index b9f477cb6d..0ab8c6b0bb 100644
--- a/src/backends/reference/workloads/CMakeLists.txt
+++ b/src/backends/reference/workloads/CMakeLists.txt
@@ -68,6 +68,9 @@ list(APPEND armnnRefBackendWorkloads_sources
     RefBatchToSpaceNdWorkload.hpp
     RefCastWorkload.cpp
     RefCastWorkload.hpp
+    RefChannelShuffleWorkload.cpp
+    RefChannelShuffleWorkload.hpp
+    RefShapeWorkload.hpp
     RefComparisonWorkload.cpp
     RefComparisonWorkload.hpp
     RefConcatWorkload.cpp
diff --git a/src/backends/reference/workloads/RefChannelShuffleWorkload.cpp b/src/backends/reference/workloads/RefChannelShuffleWorkload.cpp
new file mode 100644
index 0000000000..6571715c63
--- /dev/null
+++ b/src/backends/reference/workloads/RefChannelShuffleWorkload.cpp
@@ -0,0 +1,83 @@
+//
+// Copyright © 2021 Arm Ltd and Contributors. All rights reserved.
+// SPDX-License-Identifier: MIT
+//
+
+#include <backendsCommon/test/DataTypeUtils.hpp>
+#include <armnn/backends/ITensorHandleFactory.hpp>
+#include <armnnUtils/Transpose.hpp>
+#include "RefChannelShuffleWorkload.hpp"
+#include "RefWorkloadUtils.hpp"
+#include "Profiling.hpp"
+#include "Decoders.hpp"
+#include "Encoders.hpp"
+
+namespace armnn
+{
+void RefChannelShuffleWorkload::Execute() const
+{
+    Execute(m_Data.m_Inputs, m_Data.m_Outputs);
+}
+
+void RefChannelShuffleWorkload::ExecuteAsync(WorkingMemDescriptor &workingMemDescriptor)
+{
+    Execute(workingMemDescriptor.m_Inputs, workingMemDescriptor.m_Outputs);
+}
+
+// Reference implementation for channel shuffle taken from
+// https://android.googlesource.com/platform/frameworks/ml/+/refs/heads/master/nn/common/operations/ChannelShuffle.cpp
+void RefChannelShuffleWorkload::Execute(std::vector<ITensorHandle*> inputs,
+                                        std::vector<ITensorHandle*> outputs) const
+{
+    ARMNN_SCOPED_PROFILING_EVENT(Compute::CpuRef, "RefChannelShuffleWorkload_Execute");
+
+    const TensorInfo& inputInfo  = GetTensorInfo(inputs[0]);
+    const TensorInfo& outputInfo = GetTensorInfo(outputs[0]);
+    std::unique_ptr<Decoder<float>> decoderPtr = MakeDecoder<float>(inputInfo, inputs[0]->Map());
+    Decoder<float>& decoder = *decoderPtr;
+
+    std::unique_ptr<Encoder<float>> encoderPtr = MakeEncoder<float>(outputInfo, outputs[0]->Map());
+    Encoder<float>& encoder = *encoderPtr;
+
+    auto getNumberOfElements = [](const TensorShape& tensorShape,uint32_t startAxis, uint32_t lastAxis)
+    {
+        uint32_t count = 1;
+        for (uint32_t i = startAxis; i < lastAxis; i++)
+        {
+            count *= tensorShape[i];
+        }
+        return count;
+    };
+    const TensorShape tensorShape = GetTensorInfo(inputs[0]).GetShape();
+    uint32_t channelsAxis = m_Data.m_Parameters.m_Axis; // channelsAxis to perform channel shuffle on
+
+    const uint32_t numGroups = m_Data.m_Parameters.m_NumGroups;
+    const uint32_t groupSize = tensorShape[channelsAxis] / numGroups;
+
+    uint32_t outerSize = getNumberOfElements(tensorShape, 0, channelsAxis);
+    uint32_t innerSize = getNumberOfElements(tensorShape, channelsAxis + 1, tensorShape.GetNumDimensions());
+
+    for (uint32_t outer = 0; outer < outerSize; ++outer)
+    {
+        for (uint32_t inner = 0; inner < innerSize; ++inner)
+        {
+            uint32_t decoderStep1 = outer * tensorShape[channelsAxis] * innerSize + inner;
+            decoder += decoderStep1;
+            uint32_t encoderStep1 = outer * tensorShape[channelsAxis] * innerSize + inner;
+            encoder += encoderStep1;
+            for (uint32_t i = 0; i < groupSize; i++)
+            {
+                for (uint32_t j = 0; j < numGroups; j++, encoder += innerSize, encoderStep1 += innerSize)
+                {
+                    decoder += innerSize * (i + j * groupSize);
+                    float decoded = decoder.Get();
+                    encoder.Set(decoded);
+                    decoder -= innerSize * (i + j * groupSize);
+                }
+            }
+            decoder -= decoderStep1;
+            encoder -= encoderStep1;
+        }
+    }
+}
+}
\ No newline at end of file
diff --git a/src/backends/reference/workloads/RefChannelShuffleWorkload.hpp b/src/backends/reference/workloads/RefChannelShuffleWorkload.hpp
new file mode 100644
index 0000000000..0e4c4546fc
--- /dev/null
+++ b/src/backends/reference/workloads/RefChannelShuffleWorkload.hpp
@@ -0,0 +1,25 @@
+//
+// Copyright © 2021 Arm Ltd and Contributors. All rights reserved.
+// SPDX-License-Identifier: MIT
+//
+
+#pragma once
+
+#include <backendsCommon/Workload.hpp>
+#include <backendsCommon/WorkloadData.hpp>
+
+namespace armnn
+{
+
+class RefChannelShuffleWorkload : public BaseWorkload<ChannelShuffleQueueDescriptor>
+{
+public:
+    using BaseWorkload<ChannelShuffleQueueDescriptor>::BaseWorkload;
+    void Execute() const override;
+    void ExecuteAsync(WorkingMemDescriptor& workingMemDescriptor) override;
+
+private:
+    void Execute(std::vector<ITensorHandle*> inputs, std::vector<ITensorHandle*> outputs) const;
+};
+
+} // namespace armnn
\ No newline at end of file
diff --git a/src/backends/reference/workloads/RefWorkloads.hpp b/src/backends/reference/workloads/RefWorkloads.hpp
index d3ae58ea15..1cf84eed9e 100644
--- a/src/backends/reference/workloads/RefWorkloads.hpp
+++ b/src/backends/reference/workloads/RefWorkloads.hpp
@@ -19,6 +19,7 @@
 #include "RefBatchNormalizationWorkload.hpp"
 #include "RefBatchToSpaceNdWorkload.hpp"
 #include "RefCastWorkload.hpp"
+#include "RefChannelShuffleWorkload.hpp"
 #include "RefComparisonWorkload.hpp"
 #include "RefConvolution2dWorkload.hpp"
 #include "RefConstantWorkload.hpp"