1 //
2 // Copyright © 2017 Arm Ltd. All rights reserved.
3 // SPDX-License-Identifier: MIT
4 //
5
6 #include "StridedSlice.hpp"
7
8 #include <ResolveType.hpp>
9
10 #include <armnn/utility/Assert.hpp>
11 #include <armnn/utility/NumericCast.hpp>
12
13 #include <cstring>
14
15 namespace armnn
16 {
17
18 namespace
19 {
20
PadParams(StridedSliceDescriptor & p,unsigned int dimCount)21 void PadParams(StridedSliceDescriptor& p, unsigned int dimCount)
22 {
23 ARMNN_ASSERT_MSG(dimCount <= 4, "Expected input with at most 4 dimensions");
24
25 const unsigned int beginIndicesCount =
26 armnn::numeric_cast<unsigned int>(p.m_Begin.size());
27
28 ARMNN_ASSERT(dimCount >= beginIndicesCount);
29 const unsigned int padCount = dimCount - beginIndicesCount;
30
31 p.m_Begin.resize(dimCount);
32 p.m_End.resize(dimCount);
33 p.m_Stride.resize(dimCount);
34
35 for (unsigned int i = beginIndicesCount; i > 0; --i)
36 {
37 p.m_Stride[i + padCount - 1] = p.m_Stride[i - 1];
38 p.m_Begin[i + padCount - 1] = p.m_Begin[i - 1];
39 p.m_End[i + padCount - 1] = p.m_End[i - 1];
40 }
41
42 for (unsigned int i = 0; i < padCount; ++i)
43 {
44 p.m_Stride[i] = 1;
45 p.m_Begin[i] = 0;
46 p.m_End[i] = 0;
47 }
48
49 p.m_ShrinkAxisMask <<= padCount;
50 p.m_EllipsisMask <<= padCount;
51 p.m_NewAxisMask <<= padCount;
52 p.m_BeginMask <<= padCount;
53 p.m_EndMask <<= padCount;
54 p.m_BeginMask |= (1 << padCount) - 1;
55 p.m_EndMask |= (1 << padCount) - 1;
56 }
57
LoopCondition(int index,int stop,int stride)58 bool LoopCondition(int index, int stop, int stride)
59 {
60 return stride > 0 ? index >= stop : index <= stop;
61 }
62
ExtendShape(const TensorShape & inputShape,unsigned int newNumDimensions)63 TensorShape ExtendShape(const TensorShape& inputShape,
64 unsigned int newNumDimensions)
65 {
66 if (inputShape.GetNumDimensions() >= newNumDimensions)
67 {
68 return inputShape;
69 }
70
71 std::vector<unsigned int> newSizes(newNumDimensions, 0);
72
73 unsigned int diff = newNumDimensions - inputShape.GetNumDimensions();
74
75 for (unsigned int i = 0; i < diff; i++)
76 {
77 newSizes[i] = 1;
78 }
79
80 for (unsigned int i = diff; i < newNumDimensions; i++)
81 {
82 newSizes[i] = inputShape[i - diff];
83 }
84
85 return TensorShape(newNumDimensions, newSizes.data());
86 }
87
88 } // Anonymous namespace
89
StridedSlice(const TensorInfo & inputInfo,const StridedSliceDescriptor & params,const void * inputData,void * outputData,unsigned int dataTypeSize)90 void StridedSlice(const TensorInfo& inputInfo,
91 const StridedSliceDescriptor& params,
92 const void* inputData,
93 void* outputData,
94 unsigned int dataTypeSize)
95 {
96 const unsigned char* input = reinterpret_cast<const unsigned char*>(inputData);
97 unsigned char* output = reinterpret_cast<unsigned char*>(outputData);
98
99 const TensorShape inputShape = ExtendShape(inputInfo.GetShape(), 4);
100
101 StridedSliceDescriptor paddedParams = params;
102
103 // Pad parameters to 4 dimensions
104 PadParams(paddedParams, 4);
105
106 const int start0 = paddedParams.GetStartForAxis(inputShape, 0);
107 const int stop0 = paddedParams.GetStopForAxis (inputShape, 0, start0);
108
109 const int start1 = paddedParams.GetStartForAxis(inputShape, 1);
110 const int stop1 = paddedParams.GetStopForAxis (inputShape, 1, start1);
111
112 const int start2 = paddedParams.GetStartForAxis(inputShape, 2);
113 const int stop2 = paddedParams.GetStopForAxis (inputShape, 2, start2);
114
115 const int start3 = paddedParams.GetStartForAxis(inputShape, 3);
116 const int stop3 = paddedParams.GetStopForAxis (inputShape, 3, start3);
117
118 const int step = armnn::numeric_cast<int>(dataTypeSize);
119
120 for (int in0 = start0;
121 !LoopCondition(in0, stop0, paddedParams.m_Stride[0]);
122 in0 += paddedParams.m_Stride[0])
123 {
124 for (int in1 = start1;
125 !LoopCondition(in1, stop1, paddedParams.m_Stride[1]);
126 in1 += paddedParams.m_Stride[1])
127 {
128 for (int in2 = start2;
129 !LoopCondition(in2, stop2, paddedParams.m_Stride[2]);
130 in2 += paddedParams.m_Stride[2])
131 {
132 for (int in3 = start3;
133 !LoopCondition(in3, stop3, paddedParams.m_Stride[3]);
134 in3 += paddedParams.m_Stride[3])
135 {
136 int dim1 = armnn::numeric_cast<int>(inputShape[1]);
137 int dim2 = armnn::numeric_cast<int>(inputShape[2]);
138 int dim3 = armnn::numeric_cast<int>(inputShape[3]);
139
140 int inputOffset = (((in0 * dim1 + in1) * dim2 + in2) * dim3 + in3) * step;
141 ::memcpy(output, input + inputOffset, dataTypeSize);
142 output += step;
143 }
144 }
145 }
146 }
147 }
148
149 } // namespace armnn
150