IVGCVSW-8602 Move ComputeSplitAxis() to backendsCommon/WorkloadUtils
* Use ComputeSplitAxis in SplitOperator in tosaCommon mappings
* Fix TosaRef split tests, that were missing outputInfos
Signed-off-by: Teresa Charlin <teresa.charlinreyes@arm.com>
Signed-off-by: Cathal Corbett <cathal.corbett@arm.com>
Change-Id: Ib577eacdc6399242f37d25494e208aa56db6334c
diff --git a/src/backends/backendsCommon/WorkloadUtils.cpp b/src/backends/backendsCommon/WorkloadUtils.cpp
index e36c4b2..d459820 100644
--- a/src/backends/backendsCommon/WorkloadUtils.cpp
+++ b/src/backends/backendsCommon/WorkloadUtils.cpp
@@ -1,5 +1,5 @@
//
-// Copyright © 2017-2023 Arm Ltd. All rights reserved.
+// Copyright © 2017-2024 Arm Ltd. All rights reserved.
// SPDX-License-Identifier: MIT
//
@@ -8,6 +8,7 @@
#include <armnn/Utils.hpp>
#include <armnn/utility/NumericCast.hpp>
#include <armnnUtils/DataLayoutIndexed.hpp>
+#include <armnnUtils/TensorUtils.hpp>
#include <fmt/format.h>
#include <numeric>
@@ -373,4 +374,29 @@
return permutationVector;
}
+std::set<unsigned int> ComputeSplitAxis(const armnn::SplitterDescriptor& desc, const TensorShape& input)
+{
+ unsigned int numSplit = desc.GetNumViews();
+ unsigned int numDimensions = desc.GetNumDimensions();
+ std::set<unsigned int> splitAxis;
+ if (desc.HasAxis())
+ {
+ splitAxis.insert(armnnUtils::GetUnsignedAxis(desc.GetNumDimensions(), desc.GetAxis()));
+ }
+ else
+ {
+ for (unsigned int i = 0; i < numSplit; ++i)
+ {
+ for (unsigned int dimIdx = 0; dimIdx < numDimensions; ++dimIdx)
+ {
+ if (desc.GetViewSizes(i)[dimIdx] != input[dimIdx])
+ {
+ splitAxis.insert(dimIdx);
+ }
+ }
+ }
+ }
+ return splitAxis;
+}
+
} // namespace armnn