1 //
2 // Copyright © 2020 Arm Ltd. All rights reserved.
3 // SPDX-License-Identifier: MIT
4 //
5
6 #include <armnn/Tensor.hpp>
7
8 #include <armnnUtils/Transpose.hpp>
9
10 #include "Half.hpp"
11
12 #include <cassert>
13 #include <cstring>
14
15 namespace
16 {
17
18 class TransposeLoop
19 {
20 public:
21 using size_type = unsigned int;
22
TransposeLoop(const armnn::TensorShape & srcShape,const armnn::PermutationVector & mappings)23 TransposeLoop(const armnn::TensorShape& srcShape, const armnn::PermutationVector& mappings)
24 : m_SrcShape(srcShape)
25 {
26 assert(srcShape.GetNumDimensions() == mappings.GetSize());
27
28 const size_type numDims = srcShape.GetNumDimensions();
29
30 size_type srcStride = 1U;
31 size_type dstStride = 1U;
32
33 for (size_type i = numDims - 1U, k = 0U; k < numDims; ++k, --i)
34 {
35 m_SrcStrides[i] = srcStride;
36 m_DstStrides[mappings[i]] = dstStride;
37
38 srcStride *= srcShape[i];
39 dstStride *= srcShape[mappings[i]];
40 }
41 }
42
Unroll(const void * srcData,void * dstData,size_t dataTypeSize)43 void Unroll(const void* srcData, void* dstData, size_t dataTypeSize)
44 {
45 assert(srcData);
46 assert(dstData);
47 assert(dataTypeSize > 0);
48
49 const unsigned char* srcDataPtr = reinterpret_cast<const unsigned char*>(srcData);
50 unsigned char* dstDataPtr = reinterpret_cast<unsigned char*>(dstData);
51
52 const unsigned char* const srcEndPtr = srcDataPtr + m_SrcShape.GetNumElements() * dataTypeSize;
53 unsigned char* const dstEndPtr = dstDataPtr + m_SrcShape.GetNumElements() * dataTypeSize;
54
55 Unroll(0, srcDataPtr, dstDataPtr, srcEndPtr, dstEndPtr, dataTypeSize);
56 }
57
58 private:
Unroll(size_type dimension,const unsigned char * srcData,unsigned char * dstData,const unsigned char * srcEnd,unsigned char * dstEnd,size_t dataTypeSize)59 void Unroll(size_type dimension,
60 const unsigned char* srcData, unsigned char* dstData,
61 const unsigned char* srcEnd, unsigned char* dstEnd,
62 size_t dataTypeSize)
63 {
64 assert(srcData);
65 assert(dstData);
66 assert(srcEnd);
67 assert(dstEnd);
68 assert(srcData < srcEnd);
69 assert(dstData < dstEnd);
70 assert(dataTypeSize > 0);
71
72 if (dimension >= m_SrcShape.GetNumDimensions())
73 {
74 ::memcpy(dstData, srcData, dataTypeSize);
75 }
76 else
77 {
78 for (size_type i = 0; i < m_SrcShape[dimension]; i++)
79 {
80 Unroll(dimension + 1, srcData, dstData, srcEnd, dstEnd, dataTypeSize);
81
82 srcData += m_SrcStrides[dimension] * dataTypeSize;
83 dstData += m_DstStrides[dimension] * dataTypeSize;
84 }
85 }
86 }
87
88 armnn::TensorShape m_SrcShape;
89 std::array<size_type, armnn::MaxNumOfTensorDimensions> m_SrcStrides;
90 std::array<size_type, armnn::MaxNumOfTensorDimensions> m_DstStrides;
91 };
92
93 } // namespace
94
95 namespace armnnUtils
96 {
97
TransposeTensorShape(const armnn::TensorShape & srcShape,const armnn::PermutationVector & mappings)98 armnn::TensorShape TransposeTensorShape(const armnn::TensorShape& srcShape, const armnn::PermutationVector& mappings)
99 {
100 assert(srcShape.GetNumDimensions() == mappings.GetSize());
101
102 const unsigned int numDims = mappings.GetSize();
103 unsigned int outDims[armnn::MaxNumOfTensorDimensions];
104
105 for (unsigned int i = 0U; i < numDims; ++i)
106 {
107 outDims[i] = srcShape[mappings[i]];
108 }
109 armnn::TensorShape permutedShape(numDims, outDims);
110 return permutedShape;
111 }
112
TransposeTensorShape(const armnn::TensorInfo & info,const armnn::PermutationVector & mappings)113 armnn::TensorInfo TransposeTensorShape(const armnn::TensorInfo& info, const armnn::PermutationVector& mappings)
114 {
115 armnn::TensorInfo outInfo(info);
116 outInfo.SetShape(TransposeTensorShape(info.GetShape(), mappings));
117 return outInfo;
118 }
119
Transpose(const armnn::TensorShape & srcShape,const armnn::PermutationVector & mappings,const void * src,void * dst,size_t dataTypeSize)120 void Transpose(const armnn::TensorShape& srcShape, const armnn::PermutationVector& mappings,
121 const void* src, void* dst, size_t dataTypeSize)
122 {
123 TransposeLoop(srcShape, mappings).Unroll(src, dst, dataTypeSize);
124 }
125
126 } // namespace armnnUtils
127