xref: /aosp_15_r20/external/armnn/src/backends/reference/workloads/Reduce.cpp (revision 89c4ff92f2867872bb9e2354d150bf0c8c502810)
1 //
2 // Copyright © 2021 Arm Ltd and Contributors. All rights reserved.
3 // SPDX-License-Identifier: MIT
4 //
5 
6 #include "Reduce.hpp"
7 
8 #include <armnn/utility/NumericCast.hpp>
9 
10 #include <armnn/backends/WorkloadData.hpp>
11 
12 #include <cstddef>
13 #include <functional>
14 #include <limits>
15 
16 namespace armnn
17 {
18 
NextIndex(const unsigned int numDims,const armnn::TensorShape & dims,std::vector<unsigned int> & current)19 bool NextIndex(const unsigned int numDims, const armnn::TensorShape& dims, std::vector<unsigned int>& current)
20 {
21     unsigned int carry = 1;
22 
23     for (unsigned int idx = numDims; idx-- > 0; )
24     {
25         unsigned int current_val = current[idx] + carry;
26         if (dims[idx] == current_val)
27         {
28             current[idx] = 0;
29         }
30         else
31         {
32             current[idx] = current_val;
33             carry = 0;
34             break;
35         }
36     }
37     return (carry == 0);
38 }
39 
ReducedOutputOffset(const unsigned int numDims,const armnn::TensorShape & dims,std::vector<unsigned int> & index,const unsigned int numAxis,const std::vector<unsigned int> & axis)40 unsigned int ReducedOutputOffset(const unsigned int numDims,
41                                  const armnn::TensorShape& dims,
42                                  std::vector<unsigned int>& index,
43                                  const unsigned int numAxis,
44                                  const std::vector<unsigned int>& axis)
45 {
46     unsigned int offset = 0;
47     for (unsigned int idx = 0; idx < numDims; ++idx)
48     {
49         bool isAxis = false;
50         if (!axis.empty())
51         {
52             for (unsigned int axisIdx = 0; axisIdx < numAxis; ++axisIdx)
53             {
54                 if (idx == axis[axisIdx])
55                 {
56                     isAxis = true;
57                     break;
58                 }
59             }
60         }
61         if (!isAxis)
62         {
63             offset = offset * dims[idx] + index[idx];
64         }
65     }
66     return offset;
67 }
68 
69 
Reduce(const TensorInfo & inputInfo,const TensorInfo & outputInfo,Decoder<float> & input,Encoder<float> & output,const std::vector<uint32_t> axis,const ReduceOperation reduceOperation)70 void Reduce(const TensorInfo& inputInfo,
71             const TensorInfo& outputInfo,
72             Decoder<float>& input,
73             Encoder<float>& output,
74             const std::vector<uint32_t> axis,
75             const ReduceOperation reduceOperation)
76 {
77     armnn::TensorShape inputDims = inputInfo.GetShape();
78     unsigned int inputNumDims    = inputInfo.GetNumDimensions();
79     unsigned int numOutputs      = outputInfo.GetNumElements();
80 
81     // Initialise temp output
82     std::vector<float> tempOut(numOutputs);
83     switch(reduceOperation)
84     {
85         case ReduceOperation::Mean:
86         case ReduceOperation::Sum:
87             std::fill(tempOut.begin(), tempOut.end(), 0.0f);
88             break;
89         case ReduceOperation::Prod:
90             std::fill(tempOut.begin(), tempOut.end(), 1.0f);
91             break;
92         case ReduceOperation::Max:
93             std::fill(tempOut.begin(), tempOut.end(), -1 * std::numeric_limits<float>::max());
94             break;
95         case ReduceOperation::Min:
96             std::fill(tempOut.begin(), tempOut.end(), std::numeric_limits<float>::max());
97             break;
98         default:
99             throw armnn::InvalidArgumentException("Unknown reduce method: " +
100                 std::to_string(static_cast<int>(reduceOperation)));
101     }
102 
103     // Initialise temp index
104     std::vector<unsigned int> tempIndex(inputNumDims, 0);
105 
106     std::vector<unsigned int> resolvedAxis = axis;
107     if (resolvedAxis.empty())
108     {
109         for (unsigned int idx = 0; idx < inputNumDims; ++idx)
110         {
111             resolvedAxis.push_back(idx);
112         }
113     }
114     auto numResolvedAxis = armnn::numeric_cast<unsigned int>(resolvedAxis.size());
115 
116     // Iterates through input_data and operates over the reduced axis
117     for (bool hasNext = true; hasNext; hasNext = NextIndex(inputNumDims, inputDims, tempIndex))
118     {
119         unsigned int inputOffset = ReducedOutputOffset(inputNumDims, inputDims, tempIndex, 0, {});
120         unsigned int outputOffset = ReducedOutputOffset(inputNumDims, inputDims, tempIndex,
121                                                         numResolvedAxis, resolvedAxis);
122         input[inputOffset];
123         auto inputValue = input.Get();
124         switch(reduceOperation)
125         {
126             case ReduceOperation::Mean:
127             case ReduceOperation::Sum:
128                 tempOut[outputOffset] += inputValue;
129                 break;
130             case ReduceOperation::Prod:
131                 tempOut[outputOffset] *= inputValue;
132                 break;
133             case ReduceOperation::Max:
134                 if (inputValue > tempOut[outputOffset])
135                 {
136                     tempOut[outputOffset] = inputValue;
137                 }
138                 break;
139             case ReduceOperation::Min:
140                 if (inputValue < tempOut[outputOffset])
141                 {
142                     tempOut[outputOffset] = inputValue;
143                 }
144                 break;
145             default:
146                 throw armnn::InvalidArgumentException("Unknown reduce method: " +
147                     std::to_string(static_cast<int>(reduceOperation)));
148         }
149     }
150 
151     // Takes average by num of elements added to get MEAN
152     size_t numElementsInAxis = 1;
153     for (unsigned int idx = 0; idx < numResolvedAxis; ++idx)
154     {
155         unsigned int current = inputDims[resolvedAxis[idx]];
156         ARMNN_ASSERT(armnn::numeric_cast<float>(current) <
157                      (std::numeric_limits<float>::max() / armnn::numeric_cast<float>(numElementsInAxis)));
158         numElementsInAxis *= current;
159     }
160 
161     for (unsigned int idx = 0; idx < numOutputs; ++idx)
162     {
163         output[idx];
164         if (reduceOperation == ReduceOperation::Mean)
165         {
166             if (numElementsInAxis > 0)
167             {
168                 output.Set(tempOut[idx] / armnn::numeric_cast<float>(numElementsInAxis));
169             }
170         }
171         else
172         {
173             output.Set(tempOut[idx]);
174         }
175     }
176 }
177 
178 } //namespace armnn