// // Copyright © 2017 Arm Ltd. All rights reserved. // SPDX-License-Identifier: MIT // #include "StridedSlice.hpp" #include #include 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(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 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(inputShape[1]); int dim2 = boost::numeric_cast(inputShape[2]); int dim3 = boost::numeric_cast(inputShape[3]); int inputOffset = ((in0 * dim1 + in1) * dim2 + in2) * dim3 + in3; *(outPtr++) = inputData[inputOffset]; } } } } } template void StridedSlice(const TensorInfo& inputInfo, const TensorInfo& outputInfo, const StridedSliceDescriptor& params, const float* inputData, float* outData); template void StridedSlice(const TensorInfo& inputInfo, const TensorInfo& outputInfo, const StridedSliceDescriptor& params, const uint8_t* inputData, uint8_t* outData); } //namespace armnn