1 /*
2  * Copyright (C) 2018 The Android Open Source Project
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 #define LOG_TAG "Operations"
18 
19 #include "Reduce.h"
20 
21 #include <algorithm>
22 #include <limits>
23 #include <vector>
24 
25 #include "OperationResolver.h"
26 #include "OperationsExecutionUtils.h"
27 #include "Tracing.h"
28 
29 #ifdef NN_INCLUDE_CPU_IMPLEMENTATION
30 #pragma clang diagnostic push
31 #pragma clang diagnostic ignored "-Wunused-parameter"
32 #pragma clang diagnostic ignored "-Wsign-compare"
33 #include <tensorflow/lite/kernels/internal/reference/reference_ops.h>
34 #pragma clang diagnostic pop
35 #endif  // NN_INCLUDE_CPU_IMPLEMENTATION
36 
37 namespace android {
38 namespace nn {
39 namespace reduce {
40 
41 #ifdef NN_INCLUDE_CPU_IMPLEMENTATION
42 namespace {
43 
44 template <typename T>
compute(IOperationExecutionContext * context,T init,T func (T,T))45 inline bool compute(IOperationExecutionContext* context, T init, T func(T, T)) {
46     const Shape inputShape = context->getInputShape(kInputTensor);
47     const Shape axesShape = context->getInputShape(kInputAxes);
48     const Shape outputShape = context->getOutputShape(kOutputTensor);
49     const uint32_t inputRank = getNumberOfDimensions(inputShape);
50     const uint32_t numAxes = getNumberOfElements(axesShape);
51     std::vector<int> tempIndex(inputShape.dimensions.size());
52     std::vector<int> tempAxes(numAxes);
53     return tflite::reference_ops::ReduceGeneric<T>(
54             context->getInputBuffer<T>(kInputTensor),
55             reinterpret_cast<const int32_t*>(inputShape.dimensions.data()), inputRank,
56             context->getOutputBuffer<T>(kOutputTensor),
57             reinterpret_cast<const int32_t*>(outputShape.dimensions.data()),
58             outputShape.dimensions.size(), context->getInputBuffer<int32_t>(kInputAxes), numAxes,
59             context->getInputValue<bool8>(kInputKeepDims), tempIndex.data(), tempAxes.data(), init,
60             func);
61 }
62 
63 }  // namespace
64 
prepare(IOperationExecutionContext * context)65 bool prepare(IOperationExecutionContext* context) {
66     Shape inputShape = context->getInputShape(kInputTensor);
67     const uint32_t inputRank = getNumberOfDimensions(inputShape);
68     NN_RET_CHECK_LE(inputRank, 4u);
69 
70     std::vector<bool> shouldReduce(inputRank);
71     const int32_t* axes = context->getInputBuffer<int32_t>(kInputAxes);
72     Shape axesShape = context->getInputShape(kInputAxes);
73     NN_RET_CHECK_EQ(getNumberOfDimensions(axesShape), 1u);
74     const uint32_t numAxes = getNumberOfElements(axesShape);
75     for (uint32_t i = 0; i < numAxes; ++i) {
76         int32_t axis = axes[i];
77         NN_RET_CHECK(handleNegativeAxis(inputRank, &axis));
78         shouldReduce[axis] = true;
79     }
80 
81     // Input and output must have the same quantization parameters, etc.
82     Shape outputShape = inputShape;
83     outputShape.dimensions.clear();
84     bool keepDims = context->getInputValue<bool8>(kInputKeepDims);
85     for (uint32_t axis = 0; axis < inputRank; ++axis) {
86         if (shouldReduce[axis]) {
87             if (keepDims) {
88                 outputShape.dimensions.push_back(1);
89             }
90         } else {
91             outputShape.dimensions.push_back(getSizeOfDimension(inputShape, axis));
92         }
93     }
94 
95     // Handle the case when all dimensions are removed
96     if (outputShape.dimensions.empty()) {
97         outputShape.dimensions.push_back(1);
98     }
99 
100     return context->setOutputShape(kOutputTensor, outputShape);
101 }
102 
executeProd(IOperationExecutionContext * context)103 bool executeProd(IOperationExecutionContext* context) {
104     switch (context->getInputType(kInputTensor)) {
105         case OperandType::TENSOR_FLOAT16:
106             return compute<_Float16>(context, 1, [](_Float16 a, _Float16 b) -> _Float16 {
107                 // Handle the zero case because 0 * inf evaluates to nan.
108                 if (a == 0 || b == 0) return 0;
109                 return a * b;
110             });
111         case OperandType::TENSOR_FLOAT32:
112             return compute<float>(context, 1, [](float a, float b) -> float {
113                 // Handle the zero case because 0 * inf evaluates to nan.
114                 if (a == 0 || b == 0) return 0;
115                 return a * b;
116             });
117         default:
118             NN_RET_CHECK_FAIL() << "Unsupported tensor type for operation REDUCE_PROD";
119     }
120 }
121 
executeSum(IOperationExecutionContext * context)122 bool executeSum(IOperationExecutionContext* context) {
123     switch (context->getInputType(kInputTensor)) {
124         case OperandType::TENSOR_FLOAT16:
125             return compute<_Float16>(context, 0, [](_Float16 a, _Float16 b) { return a + b; });
126         case OperandType::TENSOR_FLOAT32:
127             return compute<float>(context, 0, [](float a, float b) { return a + b; });
128         default:
129             NN_RET_CHECK_FAIL() << "Unsupported tensor type for operation REDUCE_SUM";
130     }
131 }
132 
executeMax(IOperationExecutionContext * context)133 bool executeMax(IOperationExecutionContext* context) {
134     switch (context->getInputType(kInputTensor)) {
135         case OperandType::TENSOR_FLOAT16:
136             return compute<_Float16>(context, kFloat16Lowest,
137                                      [](_Float16 a, _Float16 b) { return std::max(a, b); });
138         case OperandType::TENSOR_FLOAT32:
139             return compute<float>(context, std::numeric_limits<float>::lowest(),
140                                   [](float a, float b) { return std::max(a, b); });
141         case OperandType::TENSOR_QUANT8_ASYMM:
142             return compute<uint8_t>(context, std::numeric_limits<uint8_t>::lowest(),
143                                     [](uint8_t a, uint8_t b) { return std::max(a, b); });
144         case OperandType::TENSOR_QUANT8_ASYMM_SIGNED:
145             return compute<int8_t>(context, std::numeric_limits<int8_t>::lowest(),
146                                    [](int8_t a, int8_t b) { return std::max(a, b); });
147         default:
148             NN_RET_CHECK_FAIL() << "Unsupported tensor type for operation REDUCE_MAX";
149     }
150 }
151 
executeMin(IOperationExecutionContext * context)152 bool executeMin(IOperationExecutionContext* context) {
153     switch (context->getInputType(kInputTensor)) {
154         case OperandType::TENSOR_FLOAT16:
155             return compute<_Float16>(context, kFloat16Max,
156                                      [](_Float16 a, _Float16 b) { return std::min(a, b); });
157         case OperandType::TENSOR_FLOAT32:
158             return compute<float>(context, std::numeric_limits<float>::max(),
159                                   [](float a, float b) { return std::min(a, b); });
160         case OperandType::TENSOR_QUANT8_ASYMM:
161             return compute<uint8_t>(context, std::numeric_limits<uint8_t>::max(),
162                                     [](uint8_t a, uint8_t b) { return std::min(a, b); });
163         case OperandType::TENSOR_QUANT8_ASYMM_SIGNED:
164             return compute<int8_t>(context, std::numeric_limits<int8_t>::max(),
165                                    [](int8_t a, int8_t b) { return std::min(a, b); });
166         default:
167             NN_RET_CHECK_FAIL() << "Unsupported tensor type for operation REDUCE_MIN";
168     }
169 }
170 
executeAny(IOperationExecutionContext * context)171 bool executeAny(IOperationExecutionContext* context) {
172     switch (context->getInputType(kInputTensor)) {
173         case OperandType::TENSOR_BOOL8:
174             return compute<bool8>(context, false,
175                                   [](bool8 a, bool8 b) { return static_cast<bool8>(a || b); });
176         default:
177             NN_RET_CHECK_FAIL() << "Unsupported tensor type for operation REDUCE_ANY";
178     }
179 }
180 
executeAll(IOperationExecutionContext * context)181 bool executeAll(IOperationExecutionContext* context) {
182     switch (context->getInputType(kInputTensor)) {
183         case OperandType::TENSOR_BOOL8:
184             return compute<bool8>(context, true,
185                                   [](bool8 a, bool8 b) { return static_cast<bool8>(a && b); });
186         default:
187             NN_RET_CHECK_FAIL() << "Unsupported tensor type for operation REDUCE_ALL";
188     }
189 }
190 #endif  // NN_INCLUDE_CPU_IMPLEMENTATION
191 
192 }  // namespace reduce
193 
194 NN_REGISTER_OPERATION_DEFAULT_VALIDATION(REDUCE_PROD, reduce::prepare, reduce::executeProd);
195 NN_REGISTER_OPERATION_DEFAULT_VALIDATION(REDUCE_SUM, reduce::prepare, reduce::executeSum);
196 NN_REGISTER_OPERATION_DEFAULT_VALIDATION(REDUCE_MAX, reduce::prepare, reduce::executeMax);
197 NN_REGISTER_OPERATION_DEFAULT_VALIDATION(REDUCE_MIN, reduce::prepare, reduce::executeMin);
198 NN_REGISTER_OPERATION_DEFAULT_VALIDATION(REDUCE_ANY, reduce::prepare, reduce::executeAny);
199 NN_REGISTER_OPERATION_DEFAULT_VALIDATION(REDUCE_ALL, reduce::prepare, reduce::executeAll);
200 
201 }  // namespace nn
202 }  // namespace android
203