src/backends/reference/workloads/Mean.cpp

   1 //
   2 // Copyright © 2017 Arm Ltd. All rights reserved.
   3 // SPDX-License-Identifier: MIT
   4 //
   5
   6 #include "Mean.hpp"
   7 #include "backendsCommon/WorkloadData.hpp"
   8
   9 #include <boost/numeric/conversion/cast.hpp>
  10
  11 #include <cmath>
  12 #include <cstddef>
  13 #include <functional>
  14 #include <limits>
  15
  16 namespace armnn
  17 {
  18 bool NextIndex(const unsigned int numDims, const armnn::TensorShape& dims, std::vector<unsigned int>& current)
  19 {
  20     unsigned int carry = 1;
  21
  22     for (unsigned int idx = numDims; idx-- > 0; )
  23     {
  24         unsigned int current_val = current[idx] + carry;
  25         if (dims[idx] == current_val)
  26         {
  27             current[idx] = 0;
  28         }
  29         else
  30         {
  31             current[idx] = current_val;
  32             carry = 0;
  33             break;
  34         }
  35     }
  36     return (carry == 0);
  37 }
  38
  39 std::size_t ReducedOutputOffset(const unsigned int numDims, const armnn::TensorShape& dims,
  40                                 std::vector<unsigned int>& index, const unsigned int numAxis,
  41                                 const std::vector<unsigned int>& axis) {
  42     std::size_t offset = 0;
  43     for (unsigned int idx = 0; idx < numDims; ++idx)
  44     {
  45         bool isAxis = false;
  46         if (!axis.empty())
  47         {
  48             for (unsigned int axisIdx = 0; axisIdx < numAxis; ++axisIdx)
  49             {
  50                 if (idx == axis[axisIdx])
  51                 {
  52                     isAxis = true;
  53                     break;
  54                 }
  55             }
  56         }
  57         if (!isAxis)
  58         {
  59             offset = offset * boost::numeric_cast<size_t>(dims[idx]) + boost::numeric_cast<size_t>(index[idx]);
  60         }
  61     }
  62     return offset;
  63 }
  64 } // namespace
  65
  66 namespace armnn
  67 {
  68 void Mean(const armnn::TensorInfo& inputInfo,
  69           const armnn::TensorInfo& outputInfo,
  70           const std::vector<unsigned int>& axis,
  71           const float* inputData,
  72           float* outputData) {
  73
  74     unsigned int inputNumDims = inputInfo.GetNumDimensions();
  75     unsigned int outputNumDims = outputInfo.GetNumDimensions();
  76
  77     armnn::TensorShape outputDims = outputInfo.GetShape();
  78     armnn::TensorShape inputDims = inputInfo.GetShape();
  79
  80     // Initialise output data.
  81     size_t numOutputs = 1;
  82     for (unsigned int idx = 0; idx < outputNumDims; ++idx)
  83     {
  84         numOutputs *= boost::numeric_cast<size_t>(outputDims[idx]);
  85     }
  86
  87     std::vector<float> tempSum(numOutputs);
  88     for (size_t idx = 0; idx < numOutputs; ++idx)
  89     {
  90         outputData[idx] = 0.0f;
  91         tempSum[idx] = 0.0f;
  92     }
  93
  94     // Initialise temp index.
  95     std::vector<unsigned int> tempIndex(inputNumDims);
  96     for (unsigned int idx = 0; idx < inputNumDims; ++idx)
  97     {
  98         tempIndex[idx] = 0;
  99     }
 100
 101     std::vector<unsigned int> resolvedAxis = axis;
 102     if (resolvedAxis.empty())
 103     {
 104       for (unsigned int idx = 0; idx < inputNumDims; ++idx)
 105       {
 106           resolvedAxis.push_back(idx);
 107       }
 108     }
 109     unsigned int numResolvedAxis = boost::numeric_cast<unsigned int>(resolvedAxis.size());
 110
 111     // Iterates through input_data and sum up the reduced axis.
 112     for (bool hasNext = true; hasNext; hasNext = NextIndex(inputNumDims, inputDims, tempIndex))
 113     {
 114         size_t inputOffset = ReducedOutputOffset(inputNumDims, inputDims, tempIndex, 0, {});
 115         size_t outputOffset = ReducedOutputOffset(inputNumDims, inputDims, tempIndex,
 116                                                   numResolvedAxis, resolvedAxis);
 117         tempSum[outputOffset] += inputData[inputOffset];
 118     }
 119
 120     // Takes average by num of elements added to get mean.
 121     size_t numElementsInAxis = 1;
 122     for (unsigned int idx = 0; idx < numResolvedAxis; ++idx)
 123     {
 124         size_t current = boost::numeric_cast<size_t>(inputDims[resolvedAxis[idx]]);
 125         BOOST_ASSERT(boost::numeric_cast<float>(current) <
 126               (std::numeric_limits<float>::max() / boost::numeric_cast<float>(numElementsInAxis)));
 127         numElementsInAxis *= current;
 128     }
 129     if (numElementsInAxis > 0) {
 130         for (size_t idx = 0; idx < numOutputs; ++idx)
 131         {
 132             outputData[idx] = tempSum[idx] / boost::numeric_cast<float>(numElementsInAxis);
 133         }
 134     }
 135 }
 136 } //namespace armnn