1 files changed, 155 insertions, 0 deletions
diff --git a/compiler/mir-interpreter/src/ops/Softmax.cpp b/compiler/mir-interpreter/src/ops/Softmax.cpp
new file mode 100644
index 000000000..f263f967d
--- /dev/null
+++ b/compiler/mir-interpreter/src/ops/Softmax.cpp
@@ -0,0 +1,155 @@
+/*
+ * Copyright (c) 2019 Samsung Electronics Co., Ltd. All Rights Reserved
+ * Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+ *
+ * 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 "Softmax.h"
+#include "Common.h"
+#include "QuantizationHelpers.h"
+
+#include <mir/ShapeRange.h>
+#include <mir/Tensor.h>
+
+#include <cmath>
+
+namespace mir_interpreter
+{
+
+static inline void PopulateSoftmaxLookupTable(float *table, float input_scale, float beta)
+{
+  const float scale = -input_scale * beta;
+  const int32_t max_uint8 = std::numeric_limits<uint8_t>::max();
+  for (int32_t val = 0; val <= max_uint8; ++val)
+    table[max_uint8 - val] = expf(scale * val);
+}
+
+template <typename T> struct SoftmaxImpl
+{
+  static void run(const mir::TensorVariant &arg, int axis, mir::TensorVariant &result);
+};
+
+template <typename T>
+void SoftmaxImpl<T>::run(const mir::TensorVariant &arg, int axis, mir::TensorVariant &result)
+{
+  mir::Tensor<T> arg_accessor(arg);
+  mir::Tensor<T> res_accessor(result);
+
+  mir::Shape expsum_shape = arg.getShape();
+  expsum_shape.dim(axis) = 1;
+  mir::TensorType expsum_type(arg.getElementType(), expsum_shape);
+  mir::TensorVariant expsum(expsum_type);
+  mir::Tensor<T> expsum_accessor(expsum);
+
+  for (const auto &expsum_index : mir::ShapeRange(expsum_shape))
+  {
+    T sum = 0;
+    mir::Index arg_index = expsum_index;
+    std::int32_t axis_size = arg.getShape().dim(axis);
+    for (std::int32_t i = 0; i < axis_size; ++i)
+    {
+      arg_index.at(axis) = i;
+      sum += std::exp(arg_accessor.at(arg_index));
+    }
+    expsum_accessor.at(expsum_index) = sum;
+  }
+
+  for (const auto &res_index : mir::ShapeRange(result.getShape()))
+  {
+    mir::Index expsum_index = res_index;
+    expsum_index.at(axis) = 0;
+    res_accessor.at(res_index) =
+        std::exp(arg_accessor.at(res_index)) / expsum_accessor.at(expsum_index);
+  }
+}
+
+template <> struct SoftmaxImpl<uint8_t>
+{
+  static void run(const mir::TensorVariant &input, int axis, mir::TensorVariant &output);
+};
+
+void SoftmaxImpl<uint8_t>::run(const mir::TensorVariant &input, int axis,
+                               mir::TensorVariant &output)
+{
+  const auto &input_type = input.getType();
+  const auto &output_type = output.getType();
+
+  assert(input_type.isQuantized());
+  assert(output_type.isQuantized());
+
+  const auto input_shape = input_type.getShape();
+
+  assert(input_type.getElementType() == mir::DataType::UINT8);
+  assert(axis == input_shape.rank() - 1); // supported only last dim axis
+  (void)axis;
+
+  double input_scale = input_type.getQuantization().getScale();
+  double output_scale = output_type.getQuantization().getScale();
+
+  const int trailing_dim = input_shape.rank() - 1;
+  int excluding_last_dim = 1;
+  for (int32_t i = 0; i < input_shape.rank() - 1; i++)
+  {
+    excluding_last_dim *= input_shape.dim(i);
+  }
+  const int last_dim = input_shape.dim(trailing_dim);
+
+  const int32_t clamp_max = std::numeric_limits<uint8_t>::max();
+  const int32_t clamp_min = std::numeric_limits<uint8_t>::min();
+
+  uint8_t *input_data = reinterpret_cast<uint8_t *>(input.atOffset(0));
+
+  float table[256];
+  PopulateSoftmaxLookupTable(table, input_scale, 1.f);
+
+  uint8_t *output_data = reinterpret_cast<uint8_t *>(output.atOffset(0));
+
+  for (int i = 0; i < excluding_last_dim; ++i)
+  {
+    int32_t max_val = std::numeric_limits<uint8_t>::min();
+    // Find max quantized value.
+    for (int j = 0; j < last_dim; ++j)
+    {
+      max_val = std::max(max_val, static_cast<int32_t>(input_data[j]));
+    }
+
+    float sum_exp = 0.0f;
+    const int32_t max_uint8 = std::numeric_limits<uint8_t>::max();
+    const float *table_offset = &table[max_uint8 - max_val];
+    // Calculate normalizer sum(exp(x)).
+    for (int j = 0; j < last_dim; ++j)
+    {
+      sum_exp += table_offset[input_data[j]];
+    }
+
+    const float inv_sum_exp = 1.0f / (sum_exp * output_scale);
+    // Normalize and quantize probabilities.
+    for (int j = 0; j < last_dim; ++j)
+    {
+      const float prob_rescaled = table_offset[input_data[j]] * inv_sum_exp;
+      const int32_t prob_quantized = static_cast<int32_t>(prob_rescaled + 0.5);
+      output_data[j] =
+          static_cast<uint8_t>(std::max(std::min(clamp_max, prob_quantized), clamp_min));
+    }
+    input_data += last_dim;
+    output_data += last_dim;
+  }
+}
+
+void Softmax(const mir::TensorVariant &arg, int axis, mir::TensorVariant &result)
+{
+  dispatch<SoftmaxImpl>(arg.getElementType(), arg, axis, result);
+};
+
+} // namespace mir_interpreter