InferenceSession.cpp

#include "InferenceSession.h"
#include "graph.h"
#include "kernel_ops.h"
#include "onnx_helper.h"
#include <fstream>
#include <iostream>

bool InferenceSession::LoadModel() {
  std::ifstream ifs(model_path_, std::ios::in | std::ios::binary);
  if (!ifs) {
    std::cerr << "Failed to open ONNX file: " << model_path_ << "\n";
    return false;
  }

  google::protobuf::io::IstreamInputStream zero_copy_input(&ifs);
  google::protobuf::io::CodedInputStream coded_input(&zero_copy_input);

  if (!model_.ParseFromCodedStream(&coded_input)) {
    std::cerr << "Failed to parse ONNX ModelProto\n";
    return false;
  }

  return true;
}

bool InferenceSession::BuildGraph() {
  if (!graph_)
    graph_ = std::make_unique<Graph>();
  graph_input_names_.clear();
  graph_output_names_.clear();

  const onnx::GraphProto &g = model_.graph();
  graph_input_names_.reserve(g.input_size());
  weights_.reserve(g.initializer_size());

  for (int i = 0; i < g.initializer_size(); ++i) {
    const auto &init = g.initializer(i);
    auto &tensor = TensorFromTensorProto(init);
    weights_map_[init.name()] = std::move(tensor);
    weights_.emplace(init.name());
  }

  for (int i = 0; i < g.input_size(); ++i) {
    const onnx::ValueInfoProto &in = g.input(i);
    const std::string &name = in.name();
    if (weights_.find(name) != weights_.end())
      continue;
    graph_input_names_.push_back(name);
  }

  for (int i = 0; i < g.output_size(); ++i) {
    const onnx::ValueInfoProto &out = g.output(i);
    graph_output_names_.push_back(out.name());
  }

  for (int i = 0; i < g.node_size(); ++i) {
    const onnx::NodeProto &n = g.node(i);

    std::string node_name = n.name();
    if (node_name.empty()) {
      node_name = n.op_type() + "_" + std::to_string(i);
    }
    std::vector<std::string> inputs;
    inputs.reserve(static_cast<size_t>(n.input_size()));
    for (int j = 0; j < n.input_size(); ++j) {
      if (!n.input(j).empty())
        inputs.push_back(n.input(j));
    }

    std::vector<std::string> outputs;
    outputs.reserve(static_cast<size_t>(n.output_size()));
    for (int j = 0; j < n.output_size(); ++j) {
      if (!n.output(j).empty())
        outputs.push_back(n.output(j));
    }
    OpType op_type = StringToOpType(n.op_type());
    auto node = std::make_unique<Node>(node_name, op_type, inputs, outputs);
    graph_->addNode(std::move(node));
  }
  graph_->topoSort();
  graph_->remaining_nodes_ =
      std::make_unique<std::latch>(graph_->getSortedNodes().size());
  return true;
}

void execute(InferenceSession *session, Node *node) {

  Tensor<float> result;
  std::vector<Tensor<float>> input_tensors;
  for (const auto &input_name : node->input_names()) {
    if (session->get_weights_map().find(input_name) !=
        session->get_weights_map().end()) {
      input_tensors.push_back(session->get_weights_map().at(input_name));
    } else {
      input_tensors.push_back(session->get_values_map().at(input_name));
    }
  }
  switch (node->op_type()) {
  case OpType::Add:
    result = Add(input_tensors);
    break;
  case OpType::Gemm:
    result = Gemm(input_tensors);
    break;
  case OpType::Relu:
    result = Relu(input_tensors);
    break;
  default:
    throw std::runtime_error("Unsupported operation type");
  }

  session->get_values_map()[node->output_names()[0]] = result;
  for (auto child : node->children()) {
    int64_t old =
        child->pending_dependencies().fetch_sub(1, std::memory_order_acq_rel);
    if (old == 1) {
      // all dependencies resolved, schedule execution
      boost::asio::post(session->get_thread_pool(),
                        [session, child]() { execute(session, child); });
    }
  }
  session->get_remaining_nodes().count_down();
}

std::vector<Tensor<float>> InferenceSession::run(
    std::unordered_map<std::string, Tensor<float>> &input_tensors) {
  values_map_.clear();
  LoadModel();
  BuildGraph();

  // Initialize input input_tensors
  for (const auto &name : graph_input_names_) {
    auto it = input_tensors.find(name);
    if (it != input_tensors.end()) {
      values_map_[name] = it->second;
    } else {
      throw std::runtime_error("Missing input tensor:  + name");
    }
  }

  // traverse the constructed graph
  for (const auto &node : graph()->getSortedNodes()) {
    if (node->pending_dependencies().load(std::memory_order_acquire) == 0) {
      boost::asio::post(this->get_thread_pool(),
                        [this, node]() { execute(this, node); });
    }
  }
  this->get_remaining_nodes().wait();
  // thread_pool_.join();

  std::vector<Tensor<float>> outputs;
  for (const auto &name : graph_output_names_) {
    outputs.push_back(values_map_.at(name));
  }
  return outputs;
}