mindspore2022/mindspore/lite/tools/converter/export_model.cc

/**
 * Copyright 2021 Huawei Technologies Co., Ltd
 *
 * 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 "tools/converter/export_model.h"
#include <fstream>
#include <map>
#include <memory>
#include <set>
#include <string>
#include <vector>
#include "backend/optimizer/common/optimizer.h"
#include "include/errorcode.h"
#include "include/version.h"
#include "ir/func_graph.h"
#include "tools/anf_exporter/anf_exporter.h"
#include "tools/converter/graphdef_transform.h"
#include "tools/converter/optimizer_manager.h"
#include "tools/converter/parser/parser_utils.h"
#include "tools/optimizer/graph/control_flow_pass.h"
#include "nnacl/op_base.h"
#include "src/common/log_util.h"

namespace mindspore {
namespace lite {
namespace {
using NodesMap = std::map<std::string, std::vector<AnfNodePtr>>;
void CloneGraphInputs(const FuncGraphPtr &origin, const FuncGraphPtr &mirror, NodesMap *origin_map,
                      NodesMap *mirror_map) {
  MS_ASSERT(origin != nullptr && mirror != nullptr);
  MS_ASSERT(origin_map != nullptr && mirror_map != nullptr);
  auto origin_inputs = origin->get_inputs();
  for (auto &input : origin_inputs) {
    auto mirror_input = mirror->add_parameter();
    MS_CHECK_TRUE_RET_VOID(mirror_input != nullptr);
    if (input->abstract() != nullptr) {
      mirror_input->set_abstract(input->abstract()->Clone());
    }
    mirror_input->set_name(input->fullname_with_scope());
    MS_ASSERT(origin_map->find(input->fullname_with_scope()) != origin_map->end());
    MS_ASSERT(mirror_map->find(input->fullname_with_scope()) != mirror_map->end());
    (*origin_map)[input->fullname_with_scope()].push_back(input);
    (*mirror_map)[input->fullname_with_scope()].push_back(mirror_input);
  }
}

AnfNodePtr CloneParameterAndValueNode(const CNodePtr &cnode, size_t index, const FuncGraphPtr &mirror_graph,
                                      const converter::Flags *flags) {
  MS_ASSERT(cnode != nullptr && mirror_graph != nullptr);
  if (index >= cnode->size()) {
    MS_LOG(ERROR) << "input index out of range.";
    return nullptr;
  }
  auto node = cnode->input(index);
  if (utils::isa<mindspore::CNode>(node)) {
    MS_LOG(ERROR) << "this func cannot copy cnode.";
    return nullptr;
  }
  if (utils::isa<ValueNode>(node)) {
    auto value_node = node->cast<ValueNodePtr>();
    auto value_ptr = value_node->value();
    MS_ASSERT(value_ptr != nullptr);
    if (utils::isa<Monad>(value_ptr)) {
      std::shared_ptr<Monad> mirror_monad;
      if (utils::isa<UMonad>(value_ptr)) {
        mirror_monad = std::make_shared<UMonad>();
      } else {
        mirror_monad = std::make_shared<IOMonad>();
      }
      MS_CHECK_TRUE_RET(mirror_monad != nullptr, nullptr);
      auto monad_abs = mirror_monad->ToAbstract();
      auto mirror_value_node = NewValueNode(mirror_monad);
      MS_CHECK_TRUE_RET(mirror_value_node != nullptr, nullptr);
      mirror_value_node->set_abstract(monad_abs);
      return mirror_value_node;
    }
  }
  DataInfo data_info;
  STATUS status;
  if (utils::isa<Parameter>(node)) {
    status = FetchDataFromParameterNode(cnode, index, flags->fmk, flags->trainModel, &data_info);
  } else if (utils::isa<ValueNode>(node)) {
    status = FetchDataFromValueNode(cnode, index, flags->fmk, flags->trainModel, &data_info);
  } else {
    status = RET_ERROR;
  }
  if (status != RET_OK && status != RET_NO_CHANGE) {
    MS_LOG(ERROR) << "fetch data failed.";
    return nullptr;
  }
  if (opt::CheckPrimitiveType(cnode, prim::kPrimTupleGetItem) && data_info.data_.size() >= sizeof(int)) {
    return NewValueNode(MakeValue<int>(*reinterpret_cast<int *>(data_info.data_.data())));
  }
  ShapeVector shape_vec(data_info.shape_.begin(), data_info.shape_.end());
  if (data_info.data_type_ == kObjectTypeTensorType) {
    shape_vec = ShapeVector{static_cast<int64_t>(data_info.data_.size() / sizeof(int))};
  }
  auto tensor_info = std::make_shared<tensor::Tensor>(static_cast<TypeId>(data_info.data_type_), shape_vec);
  MS_CHECK_TRUE_RET(tensor_info != nullptr, nullptr);
  if (!data_info.data_.empty()) {
    auto tensor_data = reinterpret_cast<uint8_t *>(tensor_info->data_c());
    if (tensor_info->data().nbytes() < 0) {
      MS_LOG(ERROR) << "tensor info data size is smaller than zero.";
      return nullptr;
    }
    if (memcpy_s(tensor_data, tensor_info->data().nbytes(), data_info.data_.data(), data_info.data_.size()) != EOK) {
      MS_LOG(ERROR) << "memcpy_s failed";
      return nullptr;
    }
  }
  auto mirror_parameter = mirror_graph->add_parameter();
  MS_CHECK_TRUE_RET(mirror_parameter != nullptr, nullptr);
  if (node->abstract() != nullptr) {
    mirror_parameter->set_abstract(node->abstract()->Clone());
  }
  mirror_parameter->set_name(node->fullname_with_scope());
  mirror_parameter->set_default_param(tensor_info);
  return mirror_parameter;
}

PrimitivePtr ClonePrimitive(const CNodePtr &cnode) {
  MS_ASSERT(cnode != nullptr);
  auto origin_prim = GetValueNode<PrimitivePtr>(cnode->input(0));
  if (origin_prim == nullptr) {
    return nullptr;
  }
  PrimitivePtr prim;
  auto op_primc_fns = ops::OpPrimCRegister::GetInstance().GetPrimCMap();
  if (op_primc_fns.find(origin_prim->name()) != op_primc_fns.end()) {
    prim = op_primc_fns[origin_prim->name()]();
  } else {
    prim = std::make_shared<PrimitiveC>(origin_prim->name());
    MS_CHECK_TRUE_RET(prim != nullptr, nullptr);
    prim->set_instance_name(origin_prim->name());
  }
  prim->SetAttrs(origin_prim->attrs());
  return prim;
}
}  // namespace

FuncGraphPtr CloneFuncGraph(const FuncGraphPtr &graph, const converter::Flags *flags,
                            std::map<FuncGraphPtr, FuncGraphPtr> *cloned_func_graph) {
  MS_ASSERT(graph != nullptr);
  auto cloned_func_graph_iter = cloned_func_graph->find(graph);
  if (cloned_func_graph_iter != cloned_func_graph->end()) {
    return cloned_func_graph_iter->second;
  }
  auto mirror_graph = std::make_shared<FuncGraph>();
  MS_CHECK_TRUE_RET(mirror_graph != nullptr, nullptr);
  cloned_func_graph->insert({graph, mirror_graph});
  mirror_graph->set_attrs(graph->attrs());
  NodesMap origin_nodes;
  NodesMap mirror_nodes;
  CloneGraphInputs(graph, mirror_graph, &origin_nodes, &mirror_nodes);
  auto node_list = TopoSort(graph->get_return());
  for (auto &node : node_list) {
    if (!utils::isa<mindspore::CNode>(node)) {
      continue;
    }
    auto cnode = node->cast<CNodePtr>();
    std::vector<AnfNodePtr> node_inputs;
    size_t begin_index = 1;
    auto mirror_prim = ClonePrimitive(cnode);
    if (mirror_prim == nullptr) {
      begin_index = 0;
    }
    for (size_t i = begin_index; i < cnode->size(); ++i) {
      auto origin_input = cnode->input(i);
      MS_CHECK_TRUE_RET(origin_input != nullptr, nullptr);
      AnfNodePtr mirror_input = nullptr;
      auto value = origin_nodes[origin_input->fullname_with_scope()];
      auto iter = std::find(value.begin(), value.end(), origin_input);
      if (iter != value.end()) {
        mirror_input = mirror_nodes[origin_input->fullname_with_scope()][iter - value.begin()];
      }
      if (mirror_input == nullptr) {
        if (IsValueNode<FuncGraph>(origin_input)) {
          auto sub_func_graph = GetValueNode<FuncGraphPtr>(origin_input);
          auto mirror_sub_graph = CloneFuncGraph(sub_func_graph, flags, cloned_func_graph);
          mirror_input = NewValueNode(mirror_sub_graph);
        } else {
          mirror_input = CloneParameterAndValueNode(cnode, i, mirror_graph, flags);
        }
        if (mirror_input == nullptr) {
          MS_LOG(ERROR) << "node input cannot be found.";
          return nullptr;
        }
        origin_nodes[origin_input->fullname_with_scope()].push_back(origin_input);
        mirror_nodes[origin_input->fullname_with_scope()].push_back(mirror_input);
      }
      node_inputs.push_back(mirror_input);
    }
    auto mirror_cnode =
      mirror_prim == nullptr ? mirror_graph->NewCNode(node_inputs) : mirror_graph->NewCNode(mirror_prim, node_inputs);
    MS_CHECK_TRUE_RET(mirror_cnode != nullptr, nullptr);
    mirror_cnode->set_fullname_with_scope(cnode->fullname_with_scope());
    if (cnode->abstract() != nullptr) {
      mirror_cnode->set_abstract(cnode->abstract()->Clone());
    }
    origin_nodes[cnode->fullname_with_scope()].push_back(cnode);
    mirror_nodes[cnode->fullname_with_scope()].push_back(mirror_cnode);
    if (opt::CheckPrimitiveType(cnode, prim::kPrimReturn)) {
      mirror_graph->set_return(mirror_cnode);
    }
  }
  return mirror_graph;
}

STATUS ExportModel(const FuncGraphPtr &graph, const converter::Flags *flags) {
  CHECK_NULL_RETURN(graph);
  CHECK_NULL_RETURN(flags);
  std::map<FuncGraphPtr, FuncGraphPtr> cloned_func_graph;
  auto mirror_graph = CloneFuncGraph(graph, flags, &cloned_func_graph);
  if (mirror_graph == nullptr) {
    MS_LOG(ERROR) << "Clone funcGraph failed.";
    return RET_ERROR;
  }
  auto manager = Manage(mirror_graph, true);
  MS_CHECK_TRUE_RET(manager != nullptr, RET_ERROR);
  std::set<FuncGraphPtr> all_func_graphs;
  GetAllFuncGraph(mirror_graph, &all_func_graphs);
  for (auto &func_graph : all_func_graphs) {
    manager->AddFuncGraph(func_graph);
  }
  if (!RunOptimizerPass(mirror_graph, {"ToNHWCFormat", "InferShapePass", "SpecialNodePostProcess"})) {
    MS_LOG(ERROR) << "Run transpose opt pass failed.";
    return RET_ERROR;
  }
  auto optimizer = std::make_shared<opt::GraphOptimizer>();
  CHECK_NULL_RETURN(optimizer);
  auto graph_pm = std::make_shared<opt::PassManager>("anf graph pass manager", true);
  CHECK_NULL_RETURN(graph_pm);
  if (flags->fmk == converter::kFmkTypeTflite || flags->fmk == converter::kFmkTypeTf ||
      flags->fmk == converter::kFmkTypeOnnx) {
    graph_pm->AddPass(std::make_shared<opt::ControlFlowPass>());
  }
  optimizer->AddPassManager(graph_pm);
  if (optimizer->Optimize(mirror_graph) == nullptr) {
    MS_LOG(ERROR) << "run  graph pass failed.";
    return RET_ERROR;
  }
  auto meta_graph = Export(mirror_graph);
  if (meta_graph == nullptr) {
    MS_LOG(ERROR) << "Export to meta graph return nullptr";
    return RET_ERROR;
  }
  auto metagraph_transform = std::make_unique<GraphDefTransform>();
  if (metagraph_transform == nullptr) {
    MS_LOG(ERROR) << "Create metagraph_transform return nullptr";
    delete meta_graph;
    return RET_ERROR;
  }
  metagraph_transform->SetGraphDef(meta_graph);
  auto status = metagraph_transform->Transform(*flags);
  if (status != RET_OK) {
    MS_LOG(ERROR) << "Transform meta graph failed " << status;
    delete meta_graph;
    return RET_ERROR;
  }
  // set output tensor names to the original names, the output_names is null in nnie converter.
  auto output_names = ConverterInnerContext::GetInstance()->GetGraphOutputTensorNames();
  if (output_names.size() > meta_graph->outputIndex.size()) {
    MS_LOG(ERROR) << "the num of setting output_names is greater than actual, " << output_names.size() << " > "
                  << meta_graph->outputIndex.size() << ".";
    ReturnCode::GetSingleReturnCode()->UpdateReturnCode(RET_ERROR);
    delete meta_graph;
    return RET_ERROR;
  }
  for (size_t idx = 0; idx < output_names.size(); idx++) {
    auto &tensor = meta_graph->allTensors.at(meta_graph->outputIndex.at(idx));
    tensor->name = output_names.at(idx);
  }
  meta_graph->version = Version();
  status = MetaGraphSerializer::Save(*meta_graph, "model");
  delete meta_graph;
  std::ostringstream oss;
  if (status != RET_OK) {
    oss << "SAVE GRAPH FAILED:" << status << " " << lite::GetErrorInfo(status);
    MS_LOG(ERROR) << oss.str();
    std::cout << oss.str() << std::endl;
    return status;
  }
  return status;
}
}  // namespace lite
}  // namespace mindspore