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
diff --git a/src/backends/backendsCommon/WorkloadUtils.hpp b/src/backends/backendsCommon/WorkloadUtils.hpp
index 6350c25..0462df6 100644
--- a/src/backends/backendsCommon/WorkloadUtils.hpp
+++ b/src/backends/backendsCommon/WorkloadUtils.hpp
@@ -1,5 +1,5 @@
//
-// Copyright © 2017, 2023 Arm Ltd. All rights reserved.
+// Copyright © 2017-2024 Arm Ltd. All rights reserved.
// SPDX-License-Identifier: MIT
//
@@ -279,4 +279,11 @@
/// \return - A permutation vector that permutes the 2 last dimensions
armnn::PermutationVector GeneratePermutationVectorOnLastTwoDimensions(unsigned int rank);
+/// Calculates the axis values for split operation.
+///
+/// \param desc - Splitter Descriptor
+/// \param input - Input tensor shape
+/// \return - A set containing axis values of slitter operation
+ std::set<unsigned int> ComputeSplitAxis(const armnn::SplitterDescriptor& desc, const TensorShape& input);
+
} //namespace armnn