src/backends/reference/workloads/ArgMinMax.cpp

   1 //
   2 // Copyright © 2019 Arm Ltd. All rights reserved.
   3 // SPDX-License-Identifier: MIT
   4 //
   5
   6 #include "ArgMinMax.hpp"
   7
   8 #include <TensorUtils.hpp>
   9
  10 #include <boost/numeric/conversion/cast.hpp>
  11
  12 namespace armnn
  13 {
  14
  15 void ArgMinMax(Decoder<float>& in, int32_t* out, const TensorInfo& inputTensorInfo,
  16                const TensorInfo& outputTensorInfo, ArgMinMaxFunction function, int axis)
  17 {
  18     unsigned int uAxis = armnnUtils::GetUnsignedAxis(inputTensorInfo.GetNumDimensions(), axis);
  19
  20     const unsigned int outerElements = armnnUtils::GetNumElementsBetween(inputTensorInfo.GetShape(), 0, uAxis);
  21     const unsigned int axisSize = inputTensorInfo.GetShape()[uAxis];
  22     const unsigned int innerElements = armnnUtils::GetNumElementsBetween(inputTensorInfo.GetShape(),
  23                                                                          uAxis + 1,
  24                                                                          inputTensorInfo.GetNumDimensions());
  25
  26     for (unsigned int outer = 0; outer < outerElements; ++outer) {
  27         for (unsigned int inner = 0; inner < innerElements; ++inner) {
  28             in[outer * axisSize * innerElements + inner];
  29             auto tmpValue = in.Get();
  30             unsigned int tmpIndex = 0;
  31             for (unsigned int i = 1; i < axisSize; ++i) {
  32                 in[(outer * axisSize * innerElements) + (i * innerElements) + inner];
  33                 const auto& value = in.Get();
  34                 if ((function == armnn::ArgMinMaxFunction::Min && value < tmpValue) ||
  35                     (function == armnn::ArgMinMaxFunction::Max &&  value > tmpValue)) {
  36                     tmpValue = value;
  37                     tmpIndex = i;
  38                 }
  39             }
  40             out[outer * innerElements + inner] = boost::numeric_cast<int32_t>(tmpIndex);
  41         }
  42     }
  43 }
  44
  45 } //namespace armnn