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 "LogSoftmax.h"
20
21 #include <algorithm>
22 #include <cmath>
23 #include <vector>
24
25 #include "OperationResolver.h"
26 #include "OperationsExecutionUtils.h"
27 #include "Tracing.h"
28
29 namespace android {
30 namespace nn {
31 namespace log_softmax {
32
33 template <typename T>
compute(const T * input,const Shape & shape,T beta,uint32_t axis,T * output)34 inline bool compute(const T* input, const Shape& shape, T beta, uint32_t axis, T* output) {
35 const uint32_t outerSize = getNumberOfElements(shape, 0, axis);
36 const uint32_t axisSize = getSizeOfDimension(shape, axis);
37 const uint32_t innerSize = getNumberOfElements(shape, axis + 1, getNumberOfDimensions(shape));
38 for (uint32_t outer = 0; outer < outerSize; ++outer) {
39 for (uint32_t inner = 0; inner < innerSize; ++inner) {
40 // We subtract the maximum value from each element to ensure
41 // numerical stability, taking advantage of the following equality:
42 // exp(x[i])/sum(exp(x[i])) == exp(x[i]+C)/sum(exp(x[i]+C))
43 T maxValue = input[outer * axisSize * innerSize + inner];
44 for (uint32_t i = 1; i < axisSize; ++i) {
45 maxValue = std::max(maxValue, input[(outer * axisSize + i) * innerSize + inner]);
46 }
47
48 T sum = 0;
49 for (uint32_t i = 0; i < axisSize; ++i) {
50 sum += std::exp(static_cast<double>(
51 (input[(outer * axisSize + i) * innerSize + inner] - maxValue) * beta));
52 }
53
54 const T logSum = std::log(static_cast<double>(sum));
55 for (uint32_t i = 0; i < axisSize; ++i) {
56 output[(outer * axisSize + i) * innerSize + inner] =
57 (input[(outer * axisSize + i) * innerSize + inner] - maxValue) * beta -
58 logSum;
59 }
60 }
61 }
62 return true;
63 }
64
prepare(IOperationExecutionContext * context)65 bool prepare(IOperationExecutionContext* context) {
66 return context->setOutputShape(kOutputTensor, context->getInputShape(kInputTensor));
67 }
68
execute(IOperationExecutionContext * context)69 bool execute(IOperationExecutionContext* context) {
70 int32_t axis = context->getInputValue<int32_t>(kInputAxis);
71 NN_RET_CHECK(handleNegativeAxis(context->getInputShape(kInputTensor), &axis));
72 switch (context->getInputType(kInputTensor)) {
73 case OperandType::TENSOR_FLOAT16:
74 return compute(context->getInputBuffer<_Float16>(kInputTensor),
75 context->getInputShape(kInputTensor),
76 context->getInputValue<_Float16>(kInputBeta), axis,
77 context->getOutputBuffer<_Float16>(kOutputTensor));
78 case OperandType::TENSOR_FLOAT32:
79 return compute(context->getInputBuffer<float>(kInputTensor),
80 context->getInputShape(kInputTensor),
81 context->getInputValue<float>(kInputBeta), axis,
82 context->getOutputBuffer<float>(kOutputTensor));
83 default:
84 NN_RET_CHECK_FAIL() << "Unsupported tensor type for operation " << kOperationName;
85 }
86 }
87
88 } // namespace log_softmax
89
90 NN_REGISTER_OPERATION_DEFAULT_VALIDATION(LOG_SOFTMAX, log_softmax::prepare, log_softmax::execute);
91
92 } // namespace nn
93 } // namespace android
94