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 <algorithm>
20 #include <cmath>
21 #include <vector>
22
23 #include "OperationResolver.h"
24 #include "OperationsUtils.h"
25 #include "Tracing.h"
26
27 namespace android {
28 namespace nn {
29 namespace log_softmax {
30
31 constexpr char kOperationName[] = "LOG_SOFTMAX";
32
33 constexpr uint32_t kNumInputs = 3;
34 constexpr uint32_t kInputTensor = 0;
35 constexpr uint32_t kInputBeta = 1;
36 constexpr uint32_t kInputAxis = 2;
37
38 constexpr uint32_t kNumOutputs = 1;
39 constexpr uint32_t kOutputTensor = 0;
40
41 template <typename T>
compute(const T * input,const Shape & shape,T beta,uint32_t axis,T * output)42 inline bool compute(const T* input, const Shape& shape, T beta, uint32_t axis, T* output) {
43 const uint32_t outerSize = getNumberOfElements(shape, 0, axis);
44 const uint32_t axisSize = getSizeOfDimension(shape, axis);
45 const uint32_t innerSize = getNumberOfElements(shape, axis + 1, getNumberOfDimensions(shape));
46 for (uint32_t outer = 0; outer < outerSize; ++outer) {
47 for (uint32_t inner = 0; inner < innerSize; ++inner) {
48 // We subtract the maximum value from each element to ensure
49 // numerical stability, taking advantage of the following equality:
50 // exp(x[i])/sum(exp(x[i])) == exp(x[i]+C)/sum(exp(x[i]+C))
51 T maxValue = input[outer * axisSize * innerSize + inner];
52 for (uint32_t i = 1; i < axisSize; ++i) {
53 maxValue = std::max(maxValue, input[(outer * axisSize + i) * innerSize + inner]);
54 }
55
56 T sum = 0;
57 for (uint32_t i = 0; i < axisSize; ++i) {
58 sum += std::exp(static_cast<double>(
59 (input[(outer * axisSize + i) * innerSize + inner] - maxValue) * beta));
60 }
61
62 const T logSum = std::log(static_cast<double>(sum));
63 for (uint32_t i = 0; i < axisSize; ++i) {
64 output[(outer * axisSize + i) * innerSize + inner] =
65 (input[(outer * axisSize + i) * innerSize + inner] - maxValue) * beta -
66 logSum;
67 }
68 }
69 }
70 return true;
71 }
72
validate(const IOperationValidationContext * context)73 Result<Version> validate(const IOperationValidationContext* context) {
74 NN_RET_CHECK_EQ(context->getNumInputs(), kNumInputs);
75 NN_RET_CHECK_EQ(context->getNumOutputs(), kNumOutputs);
76 OperandType inputType = context->getInputType(kInputTensor);
77 std::vector<OperandType> inExpectedTypes;
78 std::vector<OperandType> outExpectedTypes;
79 if (inputType == OperandType::TENSOR_FLOAT32) {
80 inExpectedTypes = {OperandType::TENSOR_FLOAT32, OperandType::FLOAT32, OperandType::INT32};
81 outExpectedTypes = {OperandType::TENSOR_FLOAT32};
82 } else if (inputType == OperandType::TENSOR_FLOAT16) {
83 inExpectedTypes = {OperandType::TENSOR_FLOAT16, OperandType::FLOAT16, OperandType::INT32};
84 outExpectedTypes = {OperandType::TENSOR_FLOAT16};
85 } else {
86 return NN_ERROR() << "Unsupported input tensor type for operation " << kOperationName;
87 }
88 NN_RET_CHECK(validateInputTypes(context, inExpectedTypes));
89 NN_RET_CHECK(validateOutputTypes(context, outExpectedTypes));
90 return Version::ANDROID_Q;
91 }
92
prepare(IOperationExecutionContext * context)93 bool prepare(IOperationExecutionContext* context) {
94 return context->setOutputShape(kOutputTensor, context->getInputShape(kInputTensor));
95 }
96
execute(IOperationExecutionContext * context)97 bool execute(IOperationExecutionContext* context) {
98 int32_t axis = context->getInputValue<int32_t>(kInputAxis);
99 NN_RET_CHECK(handleNegativeAxis(context->getInputShape(kInputTensor), &axis));
100 switch (context->getInputType(kInputTensor)) {
101 case OperandType::TENSOR_FLOAT16:
102 return compute(context->getInputBuffer<_Float16>(kInputTensor),
103 context->getInputShape(kInputTensor),
104 context->getInputValue<_Float16>(kInputBeta), axis,
105 context->getOutputBuffer<_Float16>(kOutputTensor));
106 case OperandType::TENSOR_FLOAT32:
107 return compute(context->getInputBuffer<float>(kInputTensor),
108 context->getInputShape(kInputTensor),
109 context->getInputValue<float>(kInputBeta), axis,
110 context->getOutputBuffer<float>(kOutputTensor));
111 default:
112 NN_RET_CHECK_FAIL() << "Unsupported tensor type for operation " << kOperationName;
113 }
114 }
115
116 } // namespace log_softmax
117
118 NN_REGISTER_OPERATION(LOG_SOFTMAX, log_softmax::kOperationName, log_softmax::validate,
119 log_softmax::prepare, log_softmax::execute);
120
121 } // namespace nn
122 } // namespace android
123