| // |
| // Copyright © 2017 Arm Ltd. All rights reserved. |
| // SPDX-License-Identifier: MIT |
| // |
| |
| #include "StridedSlice.hpp" |
| |
| #include <boost/assert.hpp> |
| #include <boost/numeric/conversion/cast.hpp> |
| |
| namespace armnn |
| { |
| |
| void PadParams(StridedSliceDescriptor& p, unsigned int dimCount) |
| { |
| BOOST_ASSERT_MSG(dimCount <= 4, "Expected input with at most 4 dimensions"); |
| |
| const unsigned int beginIndicesCount = |
| boost::numeric_cast<unsigned int>(p.m_Begin.size()); |
| |
| BOOST_ASSERT(dimCount >= beginIndicesCount); |
| const unsigned int padCount = dimCount - beginIndicesCount; |
| |
| p.m_Begin.resize(dimCount); |
| p.m_End.resize(dimCount); |
| p.m_Stride.resize(dimCount); |
| |
| for (unsigned int i = beginIndicesCount; i > 0; --i) |
| { |
| p.m_Stride[i + padCount - 1] = p.m_Stride[i - 1]; |
| p.m_Begin[i + padCount - 1] = p.m_Begin[i - 1]; |
| p.m_End[i + padCount - 1] = p.m_End[i - 1]; |
| } |
| |
| for (unsigned int i = 0; i < padCount; ++i) |
| { |
| p.m_Stride[i] = 1; |
| p.m_Begin[i] = 0; |
| p.m_End[i] = 0; |
| } |
| |
| p.m_ShrinkAxisMask <<= padCount; |
| p.m_EllipsisMask <<= padCount; |
| p.m_NewAxisMask <<= padCount; |
| p.m_BeginMask <<= padCount; |
| p.m_EndMask <<= padCount; |
| p.m_BeginMask |= (1 << padCount) - 1; |
| p.m_EndMask |= (1 << padCount) - 1; |
| } |
| |
| bool LoopCondition(int index, int stop, int stride) |
| { |
| return stride > 0 ? index >= stop : index <= stop; |
| } |
| |
| TensorShape ExtendShape(const TensorShape& inputShape, |
| unsigned int newNumDimensions) |
| { |
| if (inputShape.GetNumDimensions() >= newNumDimensions) |
| { |
| return inputShape; |
| } |
| |
| unsigned int newSizes[newNumDimensions]; |
| |
| unsigned int diff = newNumDimensions - inputShape.GetNumDimensions(); |
| |
| for (unsigned int i = 0; i < diff; i++) |
| { |
| newSizes[i] = 1; |
| } |
| |
| for (unsigned int i = diff; i < newNumDimensions; i++) |
| { |
| newSizes[i] = inputShape[i - diff]; |
| } |
| |
| return TensorShape(newNumDimensions, newSizes); |
| } |
| |
| template<typename T> |
| void StridedSlice(const TensorInfo& inputInfo, |
| const TensorInfo& outputInfo, |
| const StridedSliceDescriptor& params, |
| const T* inputData, |
| T* outputData) |
| { |
| const TensorShape inputShape = |
| ExtendShape(inputInfo.GetShape(), 4); |
| |
| StridedSliceDescriptor paddedParams = params; |
| |
| // Pad parameters to 4 dimensions |
| PadParams(paddedParams, 4); |
| |
| const int start0 = |
| paddedParams.GetStartForAxis(inputShape, 0); |
| const int stop0 = |
| paddedParams.GetStopForAxis(inputShape, 0, start0); |
| |
| const int start1 = |
| paddedParams.GetStartForAxis(inputShape, 1); |
| const int stop1 = |
| paddedParams.GetStopForAxis(inputShape, 1, start1); |
| |
| const int start2 = |
| paddedParams.GetStartForAxis(inputShape, 2); |
| const int stop2 = |
| paddedParams.GetStopForAxis(inputShape, 2, start2); |
| |
| const int start3 = |
| paddedParams.GetStartForAxis(inputShape, 3); |
| const int stop3 = |
| paddedParams.GetStopForAxis(inputShape, 3, start3); |
| |
| T* outPtr = outputData; |
| |
| for (int in0 = start0; |
| !LoopCondition(in0, stop0, paddedParams.m_Stride[0]); |
| in0 += paddedParams.m_Stride[0]) |
| { |
| for (int in1 = start1; |
| !LoopCondition(in1, stop1, paddedParams.m_Stride[1]); |
| in1 += paddedParams.m_Stride[1]) |
| { |
| for (int in2 = start2; |
| !LoopCondition(in2, stop2, paddedParams.m_Stride[2]); |
| in2 += paddedParams.m_Stride[2]) |
| { |
| for (int in3 = start3; |
| !LoopCondition(in3, stop3, paddedParams.m_Stride[3]); |
| in3 += paddedParams.m_Stride[3]) |
| { |
| int dim1 = boost::numeric_cast<int>(inputShape[1]); |
| int dim2 = boost::numeric_cast<int>(inputShape[2]); |
| int dim3 = boost::numeric_cast<int>(inputShape[3]); |
| |
| int inputOffset = ((in0 * dim1 + in1) * dim2 + in2) * dim3 + in3; |
| *(outPtr++) = inputData[inputOffset]; |
| } |
| } |
| } |
| } |
| } |
| |
| template void StridedSlice<float>(const TensorInfo& inputInfo, |
| const TensorInfo& outputInfo, |
| const StridedSliceDescriptor& params, |
| const float* inputData, |
| float* outData); |
| |
| template void StridedSlice<uint8_t>(const TensorInfo& inputInfo, |
| const TensorInfo& outputInfo, |
| const StridedSliceDescriptor& params, |
| const uint8_t* inputData, |
| uint8_t* outData); |
| |
| } //namespace armnn |