/* * Copyright (C) 2019 The Android Open Source Project * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ #include "CpuOperationUtils.h" #include "OperationResolver.h" #include "tensorflow/lite/kernels/internal/optimized/legacy_optimized_ops.h" #include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h" #include "Tracing.h" namespace android { namespace nn { namespace softmax { constexpr char kOperationName[] = "SOFTMAX"; constexpr uint32_t kNumInputs = 3; constexpr uint32_t kInputTensor = 0; constexpr uint32_t kBetaScalar = 1; constexpr uint32_t kAxisScalar = 2; constexpr uint32_t kNumOutputs = 1; constexpr uint32_t kOutputTensor = 0; namespace { inline bool softmaxSlowFloat32(const float* inputData, const Shape& inputShape, const float beta, int32_t axis, float* outputData, const Shape& outputShape) { NNTRACE_TRANS("softmaxFloatSlow32"); const uint32_t outerSize = getNumberOfElements(inputShape, 0, axis); const uint32_t axisSize = getSizeOfDimension(inputShape, axis); const uint32_t innerSize = getNumberOfElements(inputShape, axis + 1, getNumberOfDimensions(inputShape)); for (uint32_t outer = 0; outer < outerSize; ++outer) { const float* inputBeg = inputData + outer * axisSize * innerSize; const float* inputEnd = inputBeg + axisSize * innerSize; float* outputBeg = outputData + outer * axisSize * innerSize; for (uint32_t inner = 0; inner < innerSize; ++inner, ++inputBeg, ++inputEnd, ++outputBeg) { // Find max float maxValue = -FLT_MAX; for (const float* p = inputBeg; p < inputEnd; p += innerSize) { maxValue = std::max(maxValue, *p); } // Compute sum float sum = 0.0f; for (const float* p = inputBeg; p < inputEnd; p += innerSize) { sum += std::exp((*p - maxValue) * beta); } // Compute result float* pOut = outputBeg; for (const float* p = inputBeg; p < inputEnd; p += innerSize, pOut += innerSize) { *pOut = std::exp((*p - maxValue) * beta) / sum; } } } return true; } bool softmaxFloat32(const float* inputData, const Shape& inputShape, const float beta, int32_t axis, float* outputData, const Shape& outputShape) { int32_t ndim = getNumberOfDimensions(inputShape); NN_CHECK(handleNegativeAxis(inputShape, &axis)); // TFLite optimized implementation only supports computation along the last axis if (axis == ndim - 1) { NNTRACE_COMP("optimized_ops::Softmax::float"); tflite::SoftmaxParams param = {.beta = beta}; tflite::optimized_ops::Softmax(param, convertShapeToTflshape(inputShape), inputData, convertShapeToTflshape(outputShape), outputData); return true; } else { return softmaxSlowFloat32(inputData, inputShape, beta, axis, outputData, outputShape); } } bool softmaxFloat16(const _Float16* inputData, const Shape& inputShape, const float beta, int32_t axis, _Float16* outputData, const Shape& outputShape) { NNTRACE_TRANS("softmaxFloat16"); std::vector<float> inputData_float32(getNumberOfElements(inputShape)); convertFloat16ToFloat32(inputData, &inputData_float32); std::vector<float> outputData_float32(getNumberOfElements(outputShape)); softmaxFloat32(inputData_float32.data(), inputShape, beta, axis, outputData_float32.data(), outputShape); convertFloat32ToFloat16(outputData_float32, outputData); return true; } bool softmaxQuant8Impl(const uint8_t* inputData, const Shape& inputShape, const float beta, int32_t axis, int32_t inputMultiplier, int32_t inputLeftShift, float diffMin, uint8_t* outputData, const Shape& outputShape) { NNTRACE_TRANS("softmaxQuant8"); // The representation chosen for the input to the exp() function is Q5.26. // We need to leave extra space since values that we skip might be as large as // -32 before multiplying by input_beta_multiplier, and therefore as large as // -16 afterwards. Note that exp(-8) is definitely not insignificant to // accumulation, but exp(-16) definitely is. static const int32_t kScaledDiffIntegerBits = 5; static const int kAccumulationIntegerBits = 12; using FixedPointScaledDiff = gemmlowp::FixedPoint<int32_t, kScaledDiffIntegerBits>; using FixedPointAccum = gemmlowp::FixedPoint<int32_t, kAccumulationIntegerBits>; using FixedPoint0 = gemmlowp::FixedPoint<int32_t, 0>; const uint32_t outerSize = getNumberOfElements(inputShape, 0, axis); const uint32_t axisSize = getSizeOfDimension(inputShape, axis); const uint32_t innerSize = getNumberOfElements(inputShape, axis + 1, getNumberOfDimensions(inputShape)); for (uint32_t outer = 0; outer < outerSize; ++outer) { const uint8_t* inputBeg = inputData + outer * axisSize * innerSize; const uint8_t* inputEnd = inputBeg + axisSize * innerSize; uint8_t* outputBeg = outputData + outer * axisSize * innerSize; for (uint32_t inner = 0; inner < innerSize; ++inner, ++inputBeg, ++inputEnd, ++outputBeg) { // Find max uint8_t maxValue = 0; for (const uint8_t* p = inputBeg; p < inputEnd; p += innerSize) { maxValue = std::max(maxValue, *p); } // Compute sum FixedPointAccum sum_of_exps = FixedPointAccum::Zero(); for (const uint8_t* p = inputBeg; p < inputEnd; p += innerSize) { int32_t input_diff = static_cast<int32_t>(*p) - maxValue; if (input_diff >= diffMin) { const int32_t input_diff_rescaled = tflite::MultiplyByQuantizedMultiplierGreaterThanOne( input_diff, inputMultiplier, inputLeftShift); const auto scaled_diff_f8 = FixedPointScaledDiff::FromRaw(input_diff_rescaled); sum_of_exps = sum_of_exps + gemmlowp::Rescale<kAccumulationIntegerBits>( exp_on_negative_values(scaled_diff_f8)); } } uint32_t fixed_sum_of_exps = static_cast<uint32_t>(sum_of_exps.raw()); int32_t headroom_plus_one = tflite::CountLeadingZeros(fixed_sum_of_exps); // This is the number of bits to the left of the binary point above 1.0. // Consider fixed_sum_of_exps=1.25. In that case shifted_scale=0.8 and // no later adjustment will be needed. int32_t num_bits_over_unit = kAccumulationIntegerBits - headroom_plus_one; int32_t shifted_sum_minus_one = static_cast<int32_t>( (fixed_sum_of_exps << headroom_plus_one) - (static_cast<uint32_t>(1) << 31)); FixedPoint0 shifted_scale = gemmlowp::one_over_one_plus_x_for_x_in_0_1( FixedPoint0::FromRaw(shifted_sum_minus_one)); // Compute result uint8_t* pOut = outputBeg; for (const uint8_t* p = inputBeg; p < inputEnd; p += innerSize, pOut += innerSize) { int32_t input_diff = static_cast<int32_t>(*p) - maxValue; if (input_diff >= diffMin) { const int32_t input_diff_rescaled = tflite::MultiplyByQuantizedMultiplierGreaterThanOne( input_diff, inputMultiplier, inputLeftShift); const auto scaled_diff_f8 = FixedPointScaledDiff::FromRaw(input_diff_rescaled); FixedPoint0 exp_in_0 = exp_on_negative_values(scaled_diff_f8); int32_t unsat_output = gemmlowp::RoundingDivideByPOT( (shifted_scale * exp_in_0).raw(), num_bits_over_unit + 31 - 8); *pOut = static_cast<uint8_t>( std::max(std::min(unsat_output, static_cast<int32_t>(255)), 0)); } else { *pOut = 0; } } } } return true; } bool softmaxQuant8(const uint8_t* inputData, const Shape& inputShape, const float beta, int32_t axis, uint8_t* outputData, const Shape& outputShape) { int32_t ndim = getNumberOfDimensions(inputShape); NN_CHECK(handleNegativeAxis(inputShape, &axis)); if (outputShape.offset != 0 || outputShape.scale != 1.f / 256) { LOG(ERROR) << "incorrect scale / offset for output"; return false; } static const int32_t kScaledDiffIntegerBits = 5; const double input_beta_real_multiplier = std::min(1.0 * beta * inputShape.scale * (1 << (31 - kScaledDiffIntegerBits)), (1LL << 31) - 1.0); int32_t inputMultiplier = 0, inputLeftShift = 0; if (!QuantizeMultiplierGreaterThanOne(input_beta_real_multiplier, &inputMultiplier, &inputLeftShift)) { return false; } int32_t diffMin = -CalculateInputRadius(kScaledDiffIntegerBits, inputLeftShift); // TFLite optimized implementation only supports computation along the last axis if (axis == ndim - 1) { NNTRACE_COMP("optimized_ops::Softmax::uint8"); tflite::SoftmaxParams param = {.beta = beta, .input_multiplier = inputMultiplier, .input_left_shift = inputLeftShift, .diff_min = diffMin}; tflite::optimized_ops::Softmax(param, convertShapeToTflshape(inputShape), inputData, convertShapeToTflshape(outputShape), outputData); return true; } else { return softmaxQuant8Impl(inputData, inputShape, beta, axis, inputMultiplier, inputLeftShift, diffMin, outputData, outputShape); } } } // namespace bool validate(const IOperationValidationContext* context) { NN_RET_CHECK(context->getNumInputs() == kNumInputs || context->getNumInputs() == kNumInputs - 1); NN_RET_CHECK_EQ(context->getNumOutputs(), kNumOutputs); auto inputType = context->getInputType(kInputTensor); std::vector<OperandType> inExpectedTypes; if (inputType == OperandType::TENSOR_FLOAT32 || inputType == OperandType::TENSOR_QUANT8_ASYMM) { NN_RET_CHECK(validateHalVersion(context, HalVersion::V1_0)); inExpectedTypes = {inputType, OperandType::FLOAT32}; } else if (inputType == OperandType::TENSOR_FLOAT16) { NN_RET_CHECK(validateHalVersion(context, HalVersion::V1_2)); inExpectedTypes = {inputType, OperandType::FLOAT16}; } else { NN_RET_CHECK_FAIL() << "Unsupported tensor type for operation " << kOperationName; } if (context->getNumInputs() == kNumInputs) { NN_RET_CHECK(validateHalVersion(context, HalVersion::V1_2)); inExpectedTypes.push_back(OperandType::INT32); } else { const size_t ndim = context->getInputShape(kInputTensor).dimensions.size(); if (ndim != 2 && ndim != 4 && ndim != 0) { NN_RET_CHECK(validateHalVersion(context, HalVersion::V1_2)); } } return validateInputTypes(context, inExpectedTypes) && validateOutputTypes(context, {inputType}); } bool prepare(IOperationExecutionContext* context) { Shape input = context->getInputShape(kInputTensor); float beta = (input.type == OperandType::TENSOR_FLOAT16) ? context->getInputValue<_Float16>(kBetaScalar) : context->getInputValue<float>(kBetaScalar); NN_RET_CHECK_LE(getNumberOfDimensions(input), 4); NN_RET_CHECK_GT(beta, 0.0f); Shape output = context->getOutputShape(kOutputTensor); output.dimensions = input.dimensions; return context->setOutputShape(kOutputTensor, output); } bool execute(IOperationExecutionContext* context) { // Bypass execution in the case of zero-sized input. if (getNumberOfElements(context->getOutputShape(kOutputTensor)) == 0) return true; int32_t axis = (context->getNumInputs() == kNumInputs) ? context->getInputValue<int32_t>(kAxisScalar) : -1; switch (context->getInputType(kInputTensor)) { case OperandType::TENSOR_FLOAT16: return softmaxFloat16(context->getInputBuffer<_Float16>(kInputTensor), context->getInputShape(kInputTensor), context->getInputValue<_Float16>(kBetaScalar), axis, context->getOutputBuffer<_Float16>(kOutputTensor), context->getOutputShape(kOutputTensor)); case OperandType::TENSOR_FLOAT32: return softmaxFloat32(context->getInputBuffer<float>(kInputTensor), context->getInputShape(kInputTensor), context->getInputValue<float>(kBetaScalar), axis, context->getOutputBuffer<float>(kOutputTensor), context->getOutputShape(kOutputTensor)); case OperandType::TENSOR_QUANT8_ASYMM: return softmaxQuant8(context->getInputBuffer<uint8_t>(kInputTensor), context->getInputShape(kInputTensor), context->getInputValue<float>(kBetaScalar), axis, context->getOutputBuffer<uint8_t>(kOutputTensor), context->getOutputShape(kOutputTensor)); default: NN_RET_CHECK_FAIL() << "Unsupported tensor type for operation " << kOperationName; } } } // namespace softmax NN_REGISTER_OPERATION(SOFTMAX, "SOFTMAX", softmax::validate, softmax::prepare, softmax::execute, .allowZeroSizedInput = true); } // namespace nn } // namespace android