path: root/compiler/nnkit-tf/support/src/Runner.cpp

/*
 * Copyright (c) 2019 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.
 */

// Code here refers https://github.com/Neargye/hello_tf_c_api
//
// Licensed under the MIT License <http://opensource.org/licenses/MIT>.
// SPDX-License-Identifier: MIT
// Copyright (c) 2018 - 2019 Daniil Goncharov <neargye@gmail.com>.
//
// Permission is hereby  granted, free of charge, to any  person obtaining a copy
// of this software and associated  documentation files (the "Software"), to deal
// in the Software  without restriction, including without  limitation the rights
// to  use, copy,  modify, merge,  publish, distribute,  sublicense, and/or  sell
// copies  of  the Software,  and  to  permit persons  to  whom  the Software  is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE  IS PROVIDED "AS  IS", WITHOUT WARRANTY  OF ANY KIND,  EXPRESS OR
// IMPLIED,  INCLUDING BUT  NOT  LIMITED TO  THE  WARRANTIES OF  MERCHANTABILITY,
// FITNESS FOR  A PARTICULAR PURPOSE AND  NONINFRINGEMENT. IN NO EVENT  SHALL THE
// AUTHORS  OR COPYRIGHT  HOLDERS  BE  LIABLE FOR  ANY  CLAIM,  DAMAGES OR  OTHER
// LIABILITY, WHETHER IN AN ACTION OF  CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE  OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

#include "nnkit/support/tf/Runner.h"

#include "nnkit/support/tftestinfo/ParsedTensor.h"
#include "nncc/core/ADT/tensor/Shape.h"

#include <tensorflow/c/c_api.h>

#include <vector>
#include <cassert>
#include <cstring> // std::memcpy()
#include <stdexcept>

namespace nnkit
{
namespace support
{
namespace tf
{

using nncc::core::ADT::tensor::num_elements;
using nnkit::support::tftestinfo::ParsedTensor;

namespace
{
TF_Tensor *create_tensor(const TF_DataType data_type, const std::int64_t *dims,
                         const std::size_t num_dims, const void *data, const std::size_t len)
{
  if (dims == nullptr || data == nullptr)
    return nullptr;

  TF_Tensor *tensor = TF_AllocateTensor(data_type, dims, static_cast<int>(num_dims), len);
  if (tensor == nullptr)
    return nullptr;

  void *tensor_data = TF_TensorData(tensor);
  if (tensor_data == nullptr)
  {
    TF_DeleteTensor(tensor);
    return nullptr;
  }

  std::memcpy(tensor_data, data, std::min(len, TF_TensorByteSize(tensor)));

  return tensor;
}

void deallocate_buffer(void *data, size_t)
{
  assert(data);
  std::free(data);
}

TF_Buffer *build_TFBuffer(const char *file)
{
  const auto f = std::fopen(file, "rb");

  if (f == nullptr)
    throw std::runtime_error(std::string("cannot open ") + file);

  std::fseek(f, 0, SEEK_END); // to get file size
  const auto fsize = ftell(f);

  std::fseek(f, 0, SEEK_SET);

  if (fsize < 1)
  {
    std::fclose(f);
    throw std::runtime_error(std::string("file read error:  ") + file);
  }

  const auto data = std::malloc(fsize);
  std::fread(data, fsize, 1, f);
  std::fclose(f);

  TF_Buffer *buf = TF_NewBuffer();
  buf->data = data;
  buf->length = fsize;
  buf->data_deallocator = deallocate_buffer;

  return buf;
}

} // namespace

Runner::Runner(const char *pb_path)
{
  // initialize member variables
  _sess = nullptr;
  _graph = TF_NewGraph();
  _status = TF_NewStatus();

  // import graph from file
  TF_Buffer *buffer = build_TFBuffer(pb_path);
  if (buffer == nullptr)
    throw std::runtime_error("Can't read buffer from file");

  TF_ImportGraphDefOptions *opts = TF_NewImportGraphDefOptions();

  TF_GraphImportGraphDef(_graph, buffer, opts, _status);

  TF_DeleteImportGraphDefOptions(opts);
  TF_DeleteBuffer(buffer);

  if (TF_GetCode(_status) != TF_OK) // TODO Consider wrapper to prevent memory leak
    throw std::runtime_error("Can't import GraphDef");
}

Runner::~Runner()
{
  if (_graph)
    TF_DeleteGraph(_graph);

  if (_sess)
  {
    TF_CloseSession(_sess, _status);
    TF_DeleteSession(_sess, _status);
  }

  for (auto tensor : _input_tensors)
    TF_DeleteTensor(tensor);

  for (auto tensor : _output_tensors)
    TF_DeleteTensor(tensor);

  TF_DeleteStatus(_status);
}

bool Runner::getTensorShapeFromGraphDef(const std::unique_ptr<ParsedTensor> &tensor,
                                        angkor::TensorShape &shape)
{
  assert(!tensor->hasShape());
  TF_Output tensor_op = {TF_GraphOperationByName(_graph, tensor->nodeName().c_str()),
                         tensor->tensorIndex()};

  if (tensor_op.oper == nullptr)
    return false;

  int dim_size = TF_GraphGetTensorNumDims(_graph, tensor_op, _status);
  if (dim_size == -1)
    return false;
  int64_t dims[dim_size];

  TF_GraphGetTensorShape(_graph, tensor_op, dims, dim_size, _status);

  shape.resize(dim_size);
  for (int d = 0; d < dim_size; d++)
  {
    if (dims[d] == -1)
      return false;
    shape.dim(d) = dims[d];
  }
  return true;
}

bool Runner::getTensorDtypeFromGraphDef(const std::unique_ptr<ParsedTensor> &tensor,
                                        Runner::DataType &dtype)
{
  TF_Output tensor_op = {TF_GraphOperationByName(_graph, tensor->nodeName().c_str()),
                         tensor->tensorIndex()};

  if (tensor_op.oper == nullptr)
    return false;

  TF_DataType tf_dtype = TF_OperationOutputType(tensor_op);

  switch (tf_dtype)
  {
    case TF_DataType::TF_FLOAT:
      dtype = DataType::FLOAT;
      break;
    case TF_DataType::TF_UINT8:
      dtype = DataType::U8;
      break;
    case TF_DataType::TF_UINT16:
      dtype = DataType::U16;
      break;
    case TF_DataType::TF_UINT32:
      dtype = DataType::U32;
      break;
    case TF_DataType::TF_UINT64:
      dtype = DataType::U64;
      break;
    case TF_DataType::TF_INT8:
      dtype = DataType::S8;
      break;
    case TF_DataType::TF_INT16:
      dtype = DataType::S16;
      break;
    case TF_DataType::TF_INT32:
      dtype = DataType::S32;
      break;
    case TF_DataType::TF_INT64:
      dtype = DataType::S64;
      break;
    default:
      dtype = DataType::Unknown;
      return false;
  }
  return true;
}

void Runner::prepareInputs(const std::vector<std::unique_ptr<ParsedTensor>> &inputs,
                           TensorDataMap &data_map)
{
  assert(_graph);

  for (const auto &tensor : inputs)
  {
    TF_Output input_op = {TF_GraphOperationByName(_graph, tensor->nodeName().c_str()),
                          tensor->tensorIndex()};

    if (input_op.oper == nullptr)
      throw std::runtime_error("Can't init input_op : " + tensor->name());

    std::vector<int64_t> shape;
    for (int r = 0; r < tensor->shape().rank(); r++)
      shape.emplace_back(tensor->shape().dim(r));

    int size = 0;
    if (tensor->isFloatTensor())
      size = sizeof(float);
    else
      throw std::runtime_error("Not supported tensor type");

    TF_Tensor *input_tensor =
        create_tensor(TF_FLOAT, shape.data(), shape.size(), data_map.data(tensor.get()),
                      num_elements(tensor->shape()) * size);

    _input_ops.emplace_back(input_op);
    _input_tensors.emplace_back(input_tensor);
  }
}

void Runner::prepareOutputs(const std::vector<std::unique_ptr<ParsedTensor>> &outputs)
{
  assert(_graph);

  for (const auto &tensor : outputs)
  {
    TF_Output output_op = {TF_GraphOperationByName(_graph, tensor->nodeName().c_str()),
                           tensor->tensorIndex()};

    if (output_op.oper == nullptr)
      throw std::runtime_error("Can't init output_op : " + tensor->name());

    _output_ops.emplace_back(output_op);
  }

  _output_tensors.resize(_output_ops.size());
}

void Runner::run()
{
  assert(_graph);
  assert(_output_ops.size() > 0);

  TF_SessionOptions *options = TF_NewSessionOptions();
  _sess = TF_NewSession(_graph, options, _status);
  TF_DeleteSessionOptions(options);

  if (TF_GetCode(_status) != TF_OK)
    throw std::runtime_error(TF_Message(_status));

  TF_SessionRun(_sess,
                nullptr, // Run options.
                _input_ops.data(), _input_tensors.data(), _input_ops.size(), _output_ops.data(),
                _output_tensors.data(), _output_ops.size(), nullptr,
                0,       // Target operations, number of targets.
                nullptr, // Run metadata.
                _status  // Output status.
                );

  if (TF_GetCode(_status) != TF_OK)
    throw std::runtime_error(TF_Message(_status));

  TF_CloseSession(_sess, _status);
  TF_DeleteSession(_sess, _status);
  _sess = nullptr;
}

} // namespace tf
} // namespace support
} // namespace nnkit