xref: /third_party/mindspore/mindspore/lite/src/c_api/model_c.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 "include/c_api/model_c.h"
#include <vector>
#include <cstdint>
#include "include/api/context.h"
#include "include/api/types.h"
#include "src/cxx_api/tensor/tensor_impl.h"
#include "src/cxx_api/converters.h"
#include "src/lite_session.h"

namespace mindspore {
class ModelC {
 public:
  ModelC() : session_(nullptr), context_(nullptr) {}
  ~ModelC() {
    for (auto &impl : tensor_map_) {
      delete impl.second;
    }
  }

  Status Build(const void *model_data, size_t data_size, ModelType model_type, const ContextC *model_context);
  Status Build(const std::string &model_path, ModelType model_type, const ContextC *model_context);
  Status Resize(const std::vector<MSTensor::Impl *> &inputs, const std::vector<std::vector<int64_t>> &shapes);

  Status Predict(const OH_AI_TensorHandle *inputs, size_t input_num, OH_AI_TensorHandle **outputs, size_t *output_num,
                 const OH_AI_KernelCallBack &before, const OH_AI_KernelCallBack &after);

  MSTensor::Impl **GetInputs(size_t *input_num);
  MSTensor::Impl **GetOutputs(size_t *output_num);

 private:
  std::shared_ptr<lite::LiteSession> session_ = nullptr;
  std::shared_ptr<const ContextC> context_ = nullptr;
  std::map<mindspore::tensor::MSTensor *, MSTensor::Impl *> tensor_map_;
  std::vector<MSTensor::Impl *> inputs_;
  std::vector<MSTensor::Impl *> outputs_;
  Status RunGraph(const OH_AI_KernelCallBack &before, const OH_AI_KernelCallBack &after);
  void ResetTensorData(std::vector<void *> old_data, std::vector<tensor::MSTensor *> tensors);
  MSTensor::Impl *TensorToTensorImpl(mindspore::tensor::MSTensor *tensor);
};

Status ModelC::Build(const void *model_data, size_t data_size, ModelType model_type, const ContextC *model_context) {
  if(context_.get() != model_context){
    context_.reset(model_context);
  }
  session_ = std::make_shared<lite::LiteSession>();
  if (session_ == nullptr) {
    MS_LOG(ERROR) << "create session failed";
    return kLiteNullptr;
  }
  auto ret = session_->Init(ContextUtils::Convert(model_context));
  if (ret != mindspore::lite::RET_OK) {
    MS_LOG(ERROR) << "init session failed";
    return static_cast<StatusCode>(ret);
  }
  ret = session_->LoadModelAndCompileByBuf(static_cast<const char *>(model_data), data_size);
  if (ret != RET_OK) {
    MS_LOG(ERROR) << "Load and compile failed";
  }
  return static_cast<StatusCode>(ret);
}

Status ModelC::Build(const std::string &model_path, ModelType model_type, const ContextC *model_context) {
  if(context_.get() != model_context){
    context_.reset(model_context);
  }
  session_ = std::make_shared<lite::LiteSession>();
  if (session_ == nullptr) {
    MS_LOG(ERROR) << "create session failed";
    return kLiteNullptr;
  }
  auto ret = session_->Init(ContextUtils::Convert(model_context));
  if (ret != mindspore::lite::RET_OK) {
    MS_LOG(ERROR) << "init session failed";
    return static_cast<StatusCode>(ret);
  }
  ret = session_->LoadModelAndCompileByPath(model_path);
  if (ret != RET_OK) {
    MS_LOG(ERROR) << "Load and compile failed";
  }
  return static_cast<StatusCode>(ret);
}

Status ModelC::Resize(const std::vector<MSTensor::Impl *> &inputs, const std::vector<std::vector<int64_t>> &shapes) {
  std::vector<tensor::MSTensor *> inner_input;
  size_t input_num = inputs.size();
  for (size_t i = 0; i < input_num; i++) {
    auto input = inputs[i];
    if (input == nullptr || input->lite_tensor() == nullptr) {
      MS_LOG(ERROR) << "Input tensor is null.";
      return kLiteInputTensorError;
    }
    inner_input.push_back(input->lite_tensor());
  }
  size_t shape_num = shapes.size();
  std::vector<std::vector<int32_t>> inner_shapes(shape_num);
  for (size_t i = 0; i < shape_num; i++) {
    std::transform(shapes[i].begin(), shapes[i].end(), std::back_inserter(inner_shapes[i]),
                   [](int64_t value) { return static_cast<int32_t>(value); });
  }
  if (session_ == nullptr) {
    MS_LOG(ERROR) << "Session implement is null.";
    return kLiteNullptr;
  }
  auto ret = session_->Resize(inner_input, inner_shapes);
  return static_cast<StatusCode>(ret);
}

void ModelC::ResetTensorData(std::vector<void *> old_data, std::vector<tensor::MSTensor *> tensors) {
  for (size_t j = 0; j < old_data.size(); j++) {
    tensors.at(j)->set_data(old_data.at(j));
  }
}

Status ModelC::Predict(const OH_AI_TensorHandle *inputs, size_t input_num, OH_AI_TensorHandle **outputs,
                       size_t *output_num, const OH_AI_KernelCallBack &before, const OH_AI_KernelCallBack &after) {
  if (outputs == nullptr || session_ == nullptr) {
    MS_LOG(ERROR) << "param is nullptr.";
    return kLiteError;
  }
  auto model_inputs = session_->GetInputs();
  if (model_inputs.size() != input_num) {
    MS_LOG(ERROR) << "Wrong input size.";
    return kLiteError;
  }
  std::vector<void *> old_data;
  for (size_t i = 0; i < input_num; i++) {
    auto real_input = model_inputs[i];
    auto user_input = static_cast<mindspore::MSTensor::Impl *>(inputs[i]);
    if (user_input->DataType() != static_cast<DataType>(real_input->data_type())) {
      ResetTensorData(old_data, model_inputs);
      MS_LOG(ERROR) << "DataType does not match, input:" << user_input->Name()
                    << ", real:" << real_input->tensor_name();
      return kLiteInputTensorError;
    }
    if (user_input->Data() == nullptr) {
      ResetTensorData(old_data, model_inputs);
      MS_LOG(ERROR) << "Tensor " << user_input->Name() << " has no data.";
      return kLiteInputTensorError;
    }
    old_data.push_back(real_input->data());
    if (real_input->data_type() == kObjectTypeString) {
      std::vector<int32_t> shape;
      std::transform(user_input->Shape().begin(), user_input->Shape().end(), std::back_inserter(shape),
                     [](int64_t value) { return static_cast<int32_t>(value); });
      real_input->set_shape(shape);
      real_input->set_data(user_input->MutableData());
    } else {
      if (user_input->MutableData() != real_input->data()) {
        if (real_input->Size() != user_input->DataSize()) {
          ResetTensorData(old_data, model_inputs);
          MS_LOG(ERROR) << "Tensor " << user_input->Name() << " has wrong data size.";
          return kLiteInputTensorError;
        }
        real_input->set_data(user_input->MutableData());
      }
    }
  }
  auto ret = RunGraph(before, after);
  ResetTensorData(old_data, model_inputs);
  if (ret != kSuccess) {
    MS_LOG(ERROR) << "Run graph failed.";
    return ret;
  }

  *outputs = reinterpret_cast<OH_AI_TensorHandle *>(GetOutputs(output_num));
  return kSuccess;
}

Status ModelC::RunGraph(const OH_AI_KernelCallBack &before, const OH_AI_KernelCallBack &after) {
  if (before == nullptr || after == nullptr) {
    auto ret = session_->RunGraph();
    return static_cast<StatusCode>(ret);
  }
  auto before_call_back = [&](const std::vector<mindspore::tensor::MSTensor *> &before_inputs,
                              const std::vector<mindspore::tensor::MSTensor *> &before_outputs,
                              const CallBackParam &call_param) {
    std::vector<mindspore::MSTensor::Impl> inputs_impl;
    std::vector<mindspore::MSTensor::Impl> outputs_impl;
    std::vector<OH_AI_TensorHandle> op_inputs;
    std::vector<OH_AI_TensorHandle> op_outputs;
    size_t op_input_num = before_inputs.size();
    for (size_t i = 0; i < op_input_num; i++) {
      inputs_impl.emplace_back(before_inputs[i]);
      op_inputs.push_back(&(inputs_impl.back()));
    }
    size_t op_output_num = before_outputs.size();
    for (size_t i = 0; i < op_output_num; i++) {
      outputs_impl.emplace_back(before_outputs[i]);
      op_outputs.push_back(&(outputs_impl.back()));
    }
    const OH_AI_CallBackParam op_info = {const_cast<char *>(call_param.node_name.c_str()),
                                         const_cast<char *>(call_param.node_type.c_str())};
    OH_AI_TensorHandleArray inputs = {op_input_num, op_inputs.data()};
    OH_AI_TensorHandleArray outputs = {op_output_num, op_outputs.data()};
    return before(inputs, outputs, op_info);
  };

  auto after_call_back = [&](const std::vector<mindspore::tensor::MSTensor *> &after_inputs,
                             const std::vector<mindspore::tensor::MSTensor *> &after_outputs,
                             const CallBackParam &call_param) {
    std::vector<mindspore::MSTensor::Impl> inputs_impl;
    std::vector<mindspore::MSTensor::Impl> outputs_impl;
    std::vector<OH_AI_TensorHandle> op_inputs;
    std::vector<OH_AI_TensorHandle> op_outputs;
    size_t op_input_num = after_inputs.size();
    for (size_t i = 0; i < op_input_num; i++) {
      inputs_impl.emplace_back(after_inputs[i]);
      op_inputs.push_back(&(inputs_impl.back()));
    }
    size_t op_output_num = after_outputs.size();
    for (size_t i = 0; i < op_output_num; i++) {
      outputs_impl.emplace_back(after_outputs[i]);
      op_outputs.push_back(&(outputs_impl.back()));
    }
    const OH_AI_CallBackParam op_info = {const_cast<char *>(call_param.node_name.c_str()),
                                         const_cast<char *>(call_param.node_type.c_str())};
    OH_AI_TensorHandleArray inputs = {op_input_num, op_inputs.data()};
    OH_AI_TensorHandleArray outputs = {op_output_num, op_outputs.data()};
    return after(inputs, outputs, op_info);
  };
  auto ret = session_->RunGraph(before_call_back, after_call_back);
  return static_cast<StatusCode>(ret);
}

MSTensor::Impl *ModelC::TensorToTensorImpl(mindspore::tensor::MSTensor *tensor) {
  MSTensor::Impl *impl = nullptr;
  auto iter = tensor_map_.find(tensor);
  if (iter != tensor_map_.end()) {
    impl = iter->second;
  } else {
    impl = new (std::nothrow) MSTensor::Impl(tensor);
    if (impl == nullptr || impl->lite_tensor() == nullptr) {
      MS_LOG(ERROR) << "Create tensor failed.";
      return nullptr;
    }
    tensor_map_[tensor] = impl;
  }
  return impl;
}

MSTensor::Impl **ModelC::GetInputs(size_t *input_num) {
  if (session_ == nullptr || input_num == nullptr) {
    MS_LOG(ERROR) << "Session is null.";
    return nullptr;
  }
  auto inputs = session_->GetInputs();
  *input_num = inputs.size();
  if (inputs_.capacity() < *input_num) {
    inputs_.reserve(*input_num);
  }
  inputs_.clear();
  std::transform(inputs.begin(), inputs.end(), std::back_inserter(inputs_),
                 [&](tensor::MSTensor *input) { return TensorToTensorImpl(input); });
  return inputs_.data();
}

MSTensor::Impl **ModelC::GetOutputs(size_t *output_num) {
  if (session_ == nullptr || output_num == nullptr) {
    MS_LOG(ERROR) << "Session is null.";
    return nullptr;
  }
  auto outputs = session_->GetOutputs();
  *output_num = outputs.size();
  if (outputs_.capacity() < *output_num) {
    outputs_.reserve(*output_num);
  }
  outputs_.clear();
  std::transform(outputs.begin(), outputs.end(), std::back_inserter(outputs_),
                 [&](std::unordered_map<std::string, mindspore::tensor::MSTensor *>::value_type iter) {
                   return TensorToTensorImpl(iter.second);
                 });
  return outputs_.data();
}
}  // namespace mindspore

OH_AI_ModelHandle OH_AI_ModelCreate() {
  auto impl = new (std::nothrow) mindspore::ModelC();
  if (impl == nullptr) {
    MS_LOG(ERROR) << "Model implement is null.";
    return nullptr;
  }
  return static_cast<OH_AI_ModelHandle>(impl);
}

void OH_AI_ModelDestroy(OH_AI_ModelHandle *model) {
  if (model != nullptr && *model != nullptr) {
    auto impl = static_cast<mindspore::ModelC *>(*model);
    delete impl;
    *model = nullptr;
  }
}

void OH_AI_ModelSetWorkspace(OH_AI_ModelHandle model, void *workspace, size_t workspace_size) {
  MS_LOG(ERROR) << "Unsupported Feature.";
  return;
}

OH_AI_Status OH_AI_ModelBuild(OH_AI_ModelHandle model, const void *model_data, size_t data_size,
                              OH_AI_ModelType model_type, const OH_AI_ContextHandle model_context) {
  if (model == nullptr || model_data == nullptr || model_context == nullptr) {
    MS_LOG(ERROR) << "param is nullptr.";
    return OH_AI_STATUS_LITE_NULLPTR;
  }
  if (model_type == OH_AI_MODELTYPE_INVALID) {
    MS_LOG(ERROR) << "param is invalid.";
    return OH_AI_STATUS_LITE_PARAM_INVALID;
  }
  mindspore::ContextC *context = static_cast<mindspore::ContextC *>(model_context);
  auto impl = static_cast<mindspore::ModelC *>(model);
  auto ret = impl->Build(model_data, data_size, static_cast<mindspore::ModelType>(model_type), context);
  return static_cast<OH_AI_Status>(ret.StatusCode());
}

OH_AI_Status OH_AI_ModelBuildFromFile(OH_AI_ModelHandle model, const char *model_path, OH_AI_ModelType model_type,
                                      const OH_AI_ContextHandle model_context) {
  if (model == nullptr || model_path == nullptr || model_context == nullptr) {
    MS_LOG(ERROR) << "param is nullptr.";
    return OH_AI_STATUS_LITE_NULLPTR;
  }
  if (model_type == OH_AI_MODELTYPE_INVALID) {
    MS_LOG(ERROR) << "param is invalid.";
    return OH_AI_STATUS_LITE_PARAM_INVALID;
  }
  mindspore::ContextC *context = static_cast<mindspore::ContextC *>(model_context);
  auto impl = static_cast<mindspore::ModelC *>(model);
  auto ret = impl->Build(model_path, static_cast<mindspore::ModelType>(model_type), context);
  return static_cast<OH_AI_Status>(ret.StatusCode());
}

OH_AI_Status OH_AI_ModelResize(OH_AI_ModelHandle model, const OH_AI_TensorHandleArray inputs,
                               OH_AI_ShapeInfo *shape_infos, size_t shape_info_num) {
  if (model == nullptr || shape_infos == nullptr) {
    MS_LOG(ERROR) << "param is nullptr.";
    return OH_AI_STATUS_LITE_NULLPTR;
  }
  std::vector<mindspore::MSTensor::Impl *> vec_inputs;
  std::transform(inputs.handle_list, inputs.handle_list + inputs.handle_num, std::back_inserter(vec_inputs),
                 [](OH_AI_TensorHandle value) { return static_cast<mindspore::MSTensor::Impl *>(value); });
  std::vector<std::vector<int64_t>> vec_dims;
  for (size_t i = 0; i < shape_info_num; i++) {
    std::vector<int64_t> shape(shape_infos[i].shape, shape_infos[i].shape + shape_infos[i].shape_num);
    vec_dims.push_back(shape);
  }
  auto impl = static_cast<mindspore::ModelC *>(model);
  auto ret = impl->Resize(vec_inputs, vec_dims);
  return static_cast<OH_AI_Status>(ret.StatusCode());
}

OH_AI_Status OH_AI_ModelPredict(OH_AI_ModelHandle model, const OH_AI_TensorHandleArray inputs,
                                OH_AI_TensorHandleArray *outputs, const OH_AI_KernelCallBack before,
                                const OH_AI_KernelCallBack after) {
  if (model == nullptr) {
    MS_LOG(ERROR) << "param is nullptr.";
    return OH_AI_STATUS_LITE_NULLPTR;
  }
  auto impl = static_cast<mindspore::ModelC *>(model);
  auto ret = impl->Predict(inputs.handle_list, inputs.handle_num, &(outputs->handle_list), &(outputs->handle_num),
                           before, after);
  if (!ret.IsOk()) {
    MS_LOG(ERROR) << "Predict fail, ret :" << ret;
  }
  return static_cast<OH_AI_Status>(ret.StatusCode());
}

OH_AI_TensorHandleArray OH_AI_ModelGetInputs(const OH_AI_ModelHandle model) {
  if (model == nullptr) {
    MS_LOG(ERROR) << "param is nullptr.";
    return {0, nullptr};
  }
  auto impl = static_cast<mindspore::ModelC *>(model);
  size_t input_num;
  auto handles = reinterpret_cast<OH_AI_TensorHandle *>(impl->GetInputs(&input_num));
  return {input_num, handles};
}

OH_AI_TensorHandleArray OH_AI_ModelGetOutputs(const OH_AI_ModelHandle model) {
  if (model == nullptr) {
    MS_LOG(ERROR) << "param is nullptr.";
    return {0, nullptr};
  }
  auto impl = static_cast<mindspore::ModelC *>(model);
  size_t output_num;
  auto handles = reinterpret_cast<OH_AI_TensorHandle *>(impl->GetOutputs(&output_num));
  return {output_num, handles};
}

OH_AI_TensorHandle OH_AI_ModelGetInputByTensorName(const OH_AI_ModelHandle model, const char *tensor_name) {
  if (model == nullptr || tensor_name == nullptr) {
    MS_LOG(ERROR) << "param is nullptr.";
    return nullptr;
  }
  auto impl = static_cast<mindspore::ModelC *>(model);
  size_t input_num;
  auto inputs = impl->GetInputs(&input_num);
  for (size_t i = 0; i < input_num; i++) {
    if (inputs[i]->Name() == tensor_name) {
      return static_cast<OH_AI_TensorHandle>(inputs[i]);
    }
  }
  MS_LOG(ERROR) << "tensor is not exist.";
  return nullptr;
}

OH_AI_TensorHandle OH_AI_ModelGetOutputByTensorName(const OH_AI_ModelHandle model, const char *tensor_name) {
  if (model == nullptr || tensor_name == nullptr) {
    MS_LOG(ERROR) << "param is nullptr.";
    return nullptr;
  }
  auto impl = static_cast<mindspore::ModelC *>(model);
  size_t output_num;
  auto outputs = impl->GetOutputs(&output_num);
  for (size_t i = 0; i < output_num; i++) {
    if (outputs[i]->Name() == tensor_name) {
      return static_cast<OH_AI_TensorHandle>(outputs[i]);
    }
  }
  MS_LOG(ERROR) << "tensor is not exist.";
  return nullptr;
}