tensor.h

Go to the documentation of this file.

// MIT License
//
// Copyright (c) 2020 Sergio Izquierdo
// Copyright (c) 2020 CarlPoirier
// Copyright (c) 2020 Jiannan Liu
// Copyright (c) 2020 liufeng27
// Copyright (c) 2022 Alfredo Rodriguez
//
// 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.
 
#ifndef INCLUDE_CPPFLOW_TENSOR_H_
#define INCLUDE_CPPFLOW_TENSOR_H_
 
// C headers
#include <tensorflow/c/eager/c_api.h>
#include <tensorflow/c/tf_tensor.h>
 
// C++ headers
#include <cstring>
#include <memory>
#include <string>
#include <vector>
 
// CppFlow headers
#include "cppflow/context.h"
#include "cppflow/datatype.h"
 
namespace cppflow {
 
class tensor {
 public:
  tensor() = default;
 
  template <typename T>
  tensor(const std::vector<T>& values, const std::vector<int64_t>& shape);
 
  template <typename T>
  tensor(const std::initializer_list<T>& values);
 
  template <typename T>
  tensor(const T& value);
  tensor(const tensor& tensor) = default;
  tensor(tensor&& tensor) = default;
  explicit tensor(TFE_TensorHandle* handle);
  explicit tensor(TF_Tensor* t);
 
  ~tensor() = default;
 
  tensor& operator=(const tensor& other) = default;
  tensor& operator=(tensor&& other) = default;
 
  tensor shape() const;
 
  std::string device(bool on_memory = false) const;
 
  datatype dtype() const;
 
  template <typename T>
  std::vector<T> get_data() const;
 
  // NOTE:
  //    Usually, one should not call get_eager_handle() or get_tensor() below.
  //    They are designed for implementation details in cppflow.
  //    If you are calling them directly, it is likely that you are using some
  //    tenforflow APIs not supported in cppflow.
 
  // Additional NOTE:
  //    TF_Tensor is an immutable tensor inside tensorflow.
  //    TFE_TensorHandle is a TF_Tensor and the associated device,
  //        plus some data cache
 
  // @todo Need to determine if we can mark the return value or *this as const
  std::shared_ptr<TFE_TensorHandle> get_eager_handle() const { return tfe_handle; }
 
  // Get the TF_Tensor data from the eager handle
  // Call `get_data<T>()` instead if possible
  // NOTE:
  //    Changes to the returned TF_Tensor may not be reflected in the
  //        actual device memory!
  //    Do *NOT* modify the returned TF_Tensor!
  //    See comments of `tf_tensor` for more details.
  std::shared_ptr<TF_Tensor> get_tensor() const;
 
  // DO NOT directly access this member, call get_eager_handle() instead
  // @todo This is kept as public to be compatible with existing code and
  //       should be mark as private
  std::shared_ptr<TFE_TensorHandle> tfe_handle;
 
 private:
  tensor(enum TF_DataType type, const void* data, size_t len, const std::vector<int64_t>& shape);
 
  // This member serves as a local cache of the data in tfe_handle.
  // It refers to `local_mirrors_` if on device, or `data_` if on host CPU.
  // Changes to this variable may not be reflected in the actual device memory,
  // e.g. on GPUs or on remote nodes.
  // Access it via get_tensor() if not in constructor
  mutable std::shared_ptr<TF_Tensor> tf_tensor;
};
 
}  // namespace cppflow
 
/******************************
 *   IMPLEMENTATION DETAILS   *
 ******************************/
 
namespace cppflow {
 
inline tensor::tensor(enum TF_DataType type, const void* data, size_t len, const std::vector<int64_t>& shape) {
  this->tf_tensor = {TF_AllocateTensor(type, shape.data(), static_cast<int>(shape.size()), len), TF_DeleteTensor};
  memcpy(TF_TensorData(this->tf_tensor.get()), data, TF_TensorByteSize(this->tf_tensor.get()));
  this->tfe_handle = {TFE_NewTensorHandle(this->tf_tensor.get(), context::get_status()), TFE_DeleteTensorHandle};
  status_check(context::get_status());
}
 
template <typename T>
tensor::tensor(const std::vector<T>& values, const std::vector<int64_t>& shape)
    : tensor(deduce_tf_type<T>(), values.data(), values.size() * sizeof(T), shape) {}
 
template <typename T>
tensor::tensor(const std::initializer_list<T>& values) : tensor(std::vector<T>(values), {(int64_t)values.size()}) {}
 
template <typename T>
tensor::tensor(const T& value) : tensor(std::vector<T>({value}), {}) {}
 
#ifdef TENSORFLOW_C_TF_TSTRING_H_
// For future version TensorFlow 2.4
template <>
inline tensor::tensor(const std::string& value) {
  TF_TString tstr[1];
  TF_TString_Init(&tstr[0]);
  TF_TString_Copy(&tstr[0], value.c_str(), value.size());
 
  // *this = tensor(static_cast<enum TF_DataType>(TF_STRING),
  //                reinterpret_cast<void *>(tstr), sizeof(tstr), /*shape*/ {});
  *this = tensor(static_cast<enum TF_DataType>(TF_STRING), (void*)tstr, sizeof(tstr), /*shape*/ {});
}
#else
template <>
inline tensor::tensor(const std::string& value) {
  size_t size = 8 + TF_StringEncodedSize(value.length());
  char* data = new char[value.size() + 8];
  for (int i = 0; i < 8; i++) {
    data[i] = 0;
  }
  TF_StringEncode(value.c_str(), value.size(), data + 8, size - 8, context::get_status());
  status_check(context::get_status());
 
  // *this = tensor(static_cast<enum TF_DataType>(TF_STRING),
  //                reinterpret_cast<void *>(data), size, /*shape*/ {});
  *this = tensor(static_cast<enum TF_DataType>(TF_STRING), (void*)data, size, /*shape*/ {});
  delete[] data;
}
#endif  // TENSORFLOW_C_TF_TSTRING_H_
 
inline tensor::tensor(TFE_TensorHandle* handle) { this->tfe_handle = {handle, TFE_DeleteTensorHandle}; }
 
inline tensor::tensor(TF_Tensor* t) {
  this->tf_tensor = {t, TF_DeleteTensor};
  this->tfe_handle = {TFE_NewTensorHandle(this->tf_tensor.get(), context::get_status()), TFE_DeleteTensorHandle};
  status_check(context::get_status());
}
 
inline tensor tensor::shape() const {
  auto op = TFE_NewOp(context::get_context(), "Shape", context::get_status());
  status_check(context::get_status());
 
  TFE_OpAddInput(op, this->tfe_handle.get(), context::get_status());
  status_check(context::get_status());
 
  // Output type should be int64_t
  TFE_OpSetAttrType(op, "out_type", cppflow::datatype::TF_INT64);
 
  // EXECUTE
  int n = 1;
  TFE_TensorHandle* res[1] = {nullptr};
  TFE_Execute(op, res, &n, context::get_status());
  status_check(context::get_status());
  TFE_DeleteOp(op);
 
  return tensor(res[0]);
}
 
inline std::string tensor::device(bool on_memory) const {
  std::string res;
  if (on_memory)
    res = TFE_TensorHandleBackingDeviceName(this->tfe_handle.get(), context::get_status());
  else
    res = std::string(TFE_TensorHandleDeviceName(this->tfe_handle.get(), context::get_status()));
 
  status_check(context::get_status());
  return res;
}
 
template <typename T>
std::vector<T> tensor::get_data() const {
  // Check if asked datatype and tensor datatype match
  if (this->dtype() != deduce_tf_type<T>()) {
    auto type1 = cppflow::to_string(deduce_tf_type<T>());
    auto type2 = cppflow::to_string(this->dtype());
    auto error = "Datatype in function get_data (" + type1 + ") does not match tensor datatype (" + type2 + ")";
    throw std::runtime_error(error);
  }
 
  auto res_tensor = get_tensor();
 
  // Check tensor data is not empty
  auto raw_data = TF_TensorData(res_tensor.get());
  // this->error_check(raw_data != nullptr, "Tensor data is empty");
 
  size_t size = (TF_TensorByteSize(res_tensor.get()) / TF_DataTypeSize(TF_TensorType(res_tensor.get())));
 
  // Convert to correct type
  const auto T_data = static_cast<T*>(raw_data);
  std::vector<T> r(T_data, T_data + size);
 
  return r;
}
 
inline datatype tensor::dtype() const { return TFE_TensorHandleDataType(this->tfe_handle.get()); }
 
// NOTE:
//    Changes to the returned TF_Tensor are not reflected in
//    the actual device memory!
inline std::shared_ptr<TF_Tensor> tensor::get_tensor() const {
  if (!tf_tensor) {
    tf_tensor = {TFE_TensorHandleResolve(tfe_handle.get(), context::get_status()), TF_DeleteTensor};
    status_check(context::get_status());
  }
  return tf_tensor;
}
}  // namespace cppflow
 
#endif  // INCLUDE_CPPFLOW_TENSOR_H_