src/backends/reference/workloads/StridedSlice.cpp

   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 <boost/assert.hpp>
  11 #include <boost/numeric/conversion/cast.hpp>
  12
  13 #include <cstring>
  14
  15 namespace armnn
  16 {
  17
  18 namespace
  19 {
  20
  21 void PadParams(StridedSliceDescriptor& p, unsigned int dimCount)
  22 {
  23     BOOST_ASSERT_MSG(dimCount <= 4, "Expected input with at most 4 dimensions");
  24
  25     const unsigned int beginIndicesCount =
  26         boost::numeric_cast<unsigned int>(p.m_Begin.size());
  27
  28     BOOST_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
  58 bool LoopCondition(int index, int stop, int stride)
  59 {
  60     return stride > 0 ? index >= stop : index <= stop;
  61 }
  62
  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
  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 = boost::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 = boost::numeric_cast<int>(inputShape[1]);
 137                     int dim2 = boost::numeric_cast<int>(inputShape[2]);
 138                     int dim3 = boost::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