src/backends/reference/workloads/Gather.cpp

   1 //
   2 // Copyright © 2017 Arm Ltd. All rights reserved.
   3 // SPDX-License-Identifier: MIT
   4 //
   5
   6 #include "Gather.hpp"
   7
   8 #include "RefWorkloadUtils.hpp"
   9
  10 #include <backendsCommon/WorkloadData.hpp>
  11
  12 #include <boost/core/ignore_unused.hpp>
  13 #include <boost/numeric/conversion/cast.hpp>
  14
  15 namespace armnn
  16 {
  17
  18 void Gather(const TensorInfo& paramsInfo,
  19             const TensorInfo& indicesInfo,
  20             const TensorInfo& outputInfo,
  21             Decoder<float>& params,
  22             const int32_t* indices,
  23             Encoder<float>& output)
  24 {
  25     boost::ignore_unused(outputInfo);
  26     const TensorShape& paramsShape = paramsInfo.GetShape();
  27
  28     unsigned int paramsProduct = 1;
  29     for (unsigned int i = 1; i < paramsInfo.GetNumDimensions(); ++i)
  30     {
  31         paramsProduct = paramsProduct * paramsShape[i];
  32     }
  33
  34     unsigned int outIndex = 0;
  35     for (unsigned int i = 0; i < indicesInfo.GetNumElements(); ++i)
  36     {
  37         unsigned int indx = boost::numeric_cast<unsigned int>(indices[i]);
  38
  39         BOOST_ASSERT(indices[i] >= 0 && indx < paramsShape[0]);
  40
  41         unsigned int startOffset = indx * paramsProduct;
  42         unsigned int endOffset = startOffset + paramsProduct;
  43
  44         for (unsigned int j = startOffset; j < endOffset; ++j)
  45         {
  46             params[j];
  47             float outputValue = params.Get();
  48             output[outIndex];
  49             output.Set(outputValue);
  50             ++outIndex;
  51         }
  52     }
  53
  54     BOOST_ASSERT(outIndex == outputInfo.GetNumElements());
  55 }
  56
  57 } //namespace armnn