text: /*
   * Copyright (c) 2020 Samsung Electronics Co., Ltd. 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 "luci/Pass/ResolveCustomOpMatMulPass.h"
  
  #include "flatbuffers/flexbuffers.h"
  #include <loco/IR/DataTypeTraits.h>
  
  #include <luci/IR/CircleNodes.h>
  
  #include <loco.h>
  #include <oops/InternalExn.h>
  
  namespace
  {
  
  template <typename T>
  luci::CircleConst *create_const_node(loco::Graph *g, const loco::DataType dtype,
                                       const std::vector<uint32_t> &shape,
                                       const std::vector<T> &values)
  {
    auto node = g->nodes()->create<luci::CircleConst>();
    node->dtype(dtype);
    node->rank(shape.size());
  
    uint32_t size = 1;
    for (uint32_t i = 0; i < shape.size(); ++i)
    {
      node->dim(i) = shape.at(i);
      size *= shape.at(i);
    }
    node->shape_status(luci::ShapeStatus::VALID);
  
  #define INIT_VALUES(DT)                          \
    {                                              \
      node->size<DT>(size);                        \
      for (uint32_t i = 0; i < values.size(); ++i) \
        node->at<DT>(i) = values[i];               \
    }
  
    switch (dtype)
    {
      case loco::DataType::U8:
        INIT_VALUES(loco::DataType::U8);
        break;
      case loco::DataType::S16:
        INIT_VALUES(loco::DataType::S16);
        break;
      case loco::DataType::S32:
        INIT_VALUES(loco::DataType::S32);
        break;
      case loco::DataType::FLOAT32:
        INIT_VALUES(loco::DataType::FLOAT32)
        break;
      default:
        INTERNAL_EXN("create_const_node called with unsupported type");
        break;
    }
    return node;
  }
  
  bool resolve_matmul(luci::CircleCustom *cop)
  {
  #define CHECK_OR_FALSE(condition) \
    if (not(condition))             \
      return false;
  #define CHECK_OR_THROW(condition, message) \
    if (not(condition))                      \
      INTERNAL_EXN(message);
  
    auto graph = cop->graph();
    const std::vector<uint8_t> custom_options = cop->custom_options();
    auto map = flexbuffers::GetRoot(custom_options).AsMap();
    const auto U8 = loco::DataType::U8;
    const auto S16 = loco::DataType::S16;
    const auto S32 = loco::DataType::S32;
    const auto FLOAT32 = loco::DataType::FLOAT32;
  
    bool transpose_a = map["transpose_a"].AsBool();
    bool transpose_b = map["transpose_b"].AsBool();
  
    loco::Node *lhs = cop->inputs(0);
    loco::Node *rhs = cop->inputs(1);
  
    // Check that the type of the first input is known
    auto lhs_dtype = loco::must_cast<luci::CircleNode *>(cop->inputs(0))->dtype();
    CHECK_OR_FALSE(lhs_dtype != loco::DataType::Unknown);
  
    // If transpose of first input is requested, its shape must be known
    auto circle_lhs = loco::must_cast<luci::CircleNode *>(lhs);
    CHECK_OR_FALSE(!transpose_a || circle_lhs->shape_status() == luci::ShapeStatus::VALID);
    // and its rank should be at least 2
    CHECK_OR_FALSE(!transpose_a || circle_lhs->rank() >= 2);
    // Check that the shape of the 2nd input is known
    auto circle_rhs = loco::must_cast<luci::CircleNode *>(rhs);
    CHECK_OR_FALSE(circle_rhs->shape_status() == luci::ShapeStatus::VALID);
    // TODO as of 06/23/20 TFLite only supports rank 2 for 2nd input. Fix this once that changes!
    CHECK_OR_FALSE(circle_rhs->rank() == 2);
    // Check that input data type is supported
    CHECK_OR_THROW(lhs_dtype == U8 || lhs_dtype == S16 || lhs_dtype == FLOAT32,
                   "Only UInt8, Int16 and Float32 data types are supported by MatMul");
  
    if (transpose_a)
    {
      // Create a permutation constant node
      std::vector<uint32_t> perm;
      for (uint32_t i = 0; i < circle_lhs->rank(); ++i)
        perm.push_back(i);
      std::swap(perm[circle_lhs->rank() - 1], perm[circle_lhs->rank() - 2]);
      auto perm_node = create_const_node(graph, S32, {circle_lhs->rank()}, perm);
      // Now make a transpose node
      auto transpose_node = graph->nodes()->create<luci::CircleTranspose>();
      transpose_node->a(lhs);
      transpose_node->perm(perm_node);
      lhs = transpose_node;
    }
  
    // Transpose the second input if needed. TFLite FullyConnected operator
    // assumes the second input is in column-major order, but the input is
    // in row-major order, thus we need to convert between them.
    if (!transpose_b)
    {
      const std::vector<uint32_t> perm{1, 0};
      auto perm_node = create_const_node(graph, S32, {2}, perm);
      auto transpose_node = graph->nodes()->create<luci::CircleTranspose>();
      transpose_node->a(rhs);
      transpose_node->perm(perm_node);
      rhs = transpose_node;
    }
  
    auto empty_bias = graph->nodes()->create<luci::CircleOutputExclude>();
    empty_bias->dtype(loco::DataType::FLOAT32); // Needed for type inference
  
    auto fc_node = graph->nodes()->create<luci::CircleFullyConnected>();
    fc_node->input(lhs);
    fc_node->weights(rhs);
    fc_node->bias(empty_bias);
    fc_node->fusedActivationFunction(luci::FusedActFunc::NONE);
  
    auto customOut = loco::succs(cop);
    assert(customOut.size() == 1);
    replace(*customOut.begin()).with(fc_node);
    return true;
  }
  
  } // namespace
  
  namespace luci
  {
  
  bool ResolveCustomOpMatMulPass::run(loco::Graph *g)
  {
    bool changed = false;
    for (auto node : loco::active_nodes(loco::output_nodes(g)))
    {
      auto cop = dynamic_cast<luci::CircleCustom *>(node);
      if (not cop)
        continue;
  
      if (cop->custom_code() != "MatMul")
        continue;
  
      if (!resolve_matmul(cop))
        continue;
  
      changed = true;
    }
  
    return changed;
  }
  
  } // namespace luci
  