runtime/onert/core/src/interp/operations/Softmax.cc

   1 /*
   2  * Copyright (c) 2019 Samsung Electronics Co., Ltd. All Rights Reserved
   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 #include <cker/operation/SoftMax.h>
  18
  19 #include "OperationUtil.h"
  20
  21 #include "interp/Registration.h"
  22 #include "ir/operation/Softmax.h"
  23 #include "misc/polymorphic_downcast.h"
  24
  25 namespace onert
  26 {
  27 namespace interp
  28 {
  29 namespace
  30 {
  31
  32 void Softmax2D(const float *in, const int input_size, const int batch_size, const float beta,
  33                float *out)
  34 {
  35   assert(input_size > 0);
  36
  37   // For each batch
  38   for (int b = 0; b < batch_size; b++)
  39   {
  40     // Find the max coeff.
  41     float max_coeff = in[0];
  42     for (int i = 1; i < input_size; i++)
  43     {
  44       if (in[i] > max_coeff)
  45         max_coeff = in[i];
  46     }
  47
  48     // Compute the normalized sum of exps.
  49     float exp_sum = 0.0;
  50     for (int i = 0; i < input_size; i++)
  51     {
  52       out[i] = std::exp((in[i] - max_coeff) * beta);
  53       exp_sum += out[i];
  54     }
  55
  56     // Divide by the sum of exps.
  57     float reciprocal_sum_exp = 1.f / exp_sum;
  58     for (int i = 0; i < input_size; i++)
  59     {
  60       out[i] *= reciprocal_sum_exp;
  61     }
  62
  63     // Advance in and out pointers for the next batch.
  64     in += input_size;
  65     out += input_size;
  66   }
  67 }
  68
  69 void prepareSoftMax(ExecEnv *env, const ir::Operation &node)
  70 {
  71   const auto in_index = node.getInputs().at(0);
  72   const auto out_index = node.getOutputs().at(0);
  73
  74   const auto in_tensor = env->tensorAt(in_index);
  75   UNUSED_RELEASE(in_tensor);
  76
  77   assert((in_tensor->num_dimensions() == 4) || (in_tensor->num_dimensions() == 2));
  78
  79   // Output shape should be same with input
  80   // Output type is pre-defined in model
  81   const auto output_shape = env->graph().operands().at(in_index).info().shape();
  82   const auto output_type = env->graph().operands().at(out_index).info().typeInfo();
  83
  84   const auto output_info = ir::OperandInfo::createStaticInfo(output_shape, output_type);
  85   env->allocateIfNeeded(out_index, output_info);
  86
  87   auto out_tensor = env->tensorAt(out_index);
  88   UNUSED_RELEASE(out_tensor);
  89
  90   // Check output shape is same with input
  91   assert(out_tensor->num_dimensions() == out_tensor->num_dimensions());
  92   for (uint32_t i = 0; i < in_tensor->num_dimensions(); i++)
  93   {
  94     assert(in_tensor->dimension(i) == out_tensor->dimension(i));
  95   }
  96 }
  97
  98 void invoke(const ITensor *in_tensor, const ITensor *out_tensor,
  99             const ir::operation::Softmax::Param &param)
 100 {
 101   const float *in_ptr = reinterpret_cast<const float *>(in_tensor->bufferRO());
 102   float *out_ptr = reinterpret_cast<float *>(out_tensor->buffer());
 103
 104   float beta = param.beta;
 105
 106   if (in_tensor->num_dimensions() == 2)
 107   {
 108     uint32_t batch_size = in_tensor->dimension(0);
 109     uint32_t input_size = in_tensor->dimension(1);
 110
 111     Softmax2D(in_ptr, input_size, batch_size, beta, out_ptr);
 112   }
 113   else if (in_tensor->num_dimensions() == 4)
 114   {
 115     const auto in_shape = convertShape(in_tensor->tensorInfo().shape());
 116     const auto out_shape = convertShape(out_tensor->tensorInfo().shape());
 117
 118     nnfw::cker::SoftmaxParams cker_param;
 119     cker_param.beta = beta;
 120
 121     nnfw::cker::Softmax(cker_param, in_shape, in_ptr, out_shape, out_ptr);
 122   }
 123   else
 124   {
 125     throw std::runtime_error{"Unsuported input dimension: support 2D or 4D"};
 126   }
 127 }
 128
 129 void invokeSoftMax(const ExecEnv *env, const ir::Operation &node)
 130 {
 131   const auto &softmax_node = nnfw::misc::polymorphic_downcast<const ir::operation::Softmax &>(node);
 132
 133   const auto in_index = node.getInputs().at(0);
 134   const auto out_index = node.getOutputs().at(0);
 135
 136   const auto in_tensor = env->tensorAt(in_index);
 137   const auto out_tensor = env->tensorAt(out_index);
 138
 139   const auto in_data_type = in_tensor->data_type();
 140   const auto out_data_type = out_tensor->data_type();
 141   if ((in_data_type == ir::DataType::FLOAT32) && (out_data_type == ir::DataType::FLOAT32))
 142   {
 143     invoke(in_tensor, out_tensor, softmax_node.param());
 144   }
 145   else
 146   {
 147     throw std::runtime_error{"NYI: Support float32 only"};
 148   }
 149 }
 150
 151 } // namespace
 152
 153 OpKernel *getSoftmax()
 154 {
 155   static OpKernel kernel = {prepareSoftMax, invokeSoftMax};
 156   return &kernel;
 157 }
 158
 159 } // namespace interp
 160 } // namespace onert