types.hpp

Go to the documentation of this file.

// Copyright Institute for Automotive Engineering (ika), RWTH Aachen University
// SPDX-License-Identifier: Apache-2.0
 
#pragma once
 
#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <map>
#include <memory>
#include <stdexcept>
#include <string>
#include <type_traits>
#include <variant>
#include <vector>
 
#include <Eigen/Dense>
#include <eigen3/unsupported/Eigen/CXX11/Tensor>
 
#include <model_config.pb.h>
 
#include "triton_cpp/utils.hpp"
 
namespace triton_cpp {
template <typename T>
using VectorType = typename std::conditional_t<std::is_const_v<T>,
                                               Eigen::Map<const Eigen::VectorX<typename std::remove_const_t<T>>>,
                                               Eigen::Map<Eigen::VectorX<T>>>;
 
template <typename T>
using MatrixType = typename std::conditional_t<
    std::is_const_v<T>,
    Eigen::Map<const Eigen::Matrix<typename std::remove_const_t<T>, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>>,
    Eigen::Map<Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>>>;
 
template <typename T, int rank>
using TensorType = typename std::conditional_t<
    std::is_const_v<T>,
    Eigen::TensorMap<const Eigen::Tensor<typename std::remove_const_t<T>, rank, Eigen::RowMajor, Eigen::Index>>,
    Eigen::TensorMap<Eigen::Tensor<T, rank, Eigen::RowMajor, Eigen::Index>>>;
 
using TritonDataType =
    std::variant<bool, uint8_t, uint16_t, uint32_t, uint64_t, int8_t, int16_t, int32_t, int64_t, float, double, Eigen::half>;
 
inline TritonDataType getZero(inference::DataType type) {
  switch (type) {
    case inference::DataType::TYPE_BOOL:
      return false;
    case inference::DataType::TYPE_UINT8:
      return static_cast<uint8_t>(0);
    case inference::DataType::TYPE_UINT16:
      return static_cast<uint16_t>(0);
    case inference::DataType::TYPE_UINT32:
      return static_cast<uint32_t>(0);
    case inference::DataType::TYPE_UINT64:
      return static_cast<uint64_t>(0);
    case inference::DataType::TYPE_INT8:
      return static_cast<int8_t>(0);
    case inference::DataType::TYPE_INT16:
      return static_cast<int16_t>(0);
    case inference::DataType::TYPE_INT32:
      return static_cast<int32_t>(0);
    case inference::DataType::TYPE_INT64:
      return static_cast<int64_t>(0);
    case inference::DataType::TYPE_FP16:
      return Eigen::half{0.0f};
    case inference::DataType::TYPE_FP32:
      return static_cast<float>(0);
    case inference::DataType::TYPE_FP64:
      return static_cast<double>(0);
    default:
      throw std::invalid_argument("Unsupported data type");
  }
}
 
inline std::size_t getAlignment(inference::DataType type) {
  return std::visit([](auto&& arg) { return alignof(std::decay_t<decltype(arg)>); }, getZero(type));
}
 
inline std::size_t getSharedMemoryAlignment(inference::DataType type) {
  constexpr std::size_t kMinSharedMemoryAlignment = 8;
  return std::max(getAlignment(type), kMinSharedMemoryAlignment);
}
 
inline std::size_t alignUp(std::size_t offset, std::size_t alignment) {
  if (alignment <= 1) {
    return offset;
  }
  const std::size_t remainder = offset % alignment;
  return remainder == 0 ? offset : offset + (alignment - remainder);
}
 
using ModelOutput = std::map<std::string, std::shared_ptr<triton::client::InferRequestedOutput>>;
 
struct InputData {
  std::shared_ptr<triton::client::InferInput> input;
  std::vector<uint8_t> data;
  uint8_t* data_raw{nullptr};
  uint8_t* device_data_raw{nullptr};
  std::size_t data_raw_size{0};
 
  InputData() = default;
 
  InputData(const InputData& other)
      : input{other.input},
        data{other.data},
        data_raw{other.ownsHostBuffer() ? data.data() : other.data_raw},
        device_data_raw{other.device_data_raw},
        data_raw_size{other.data_raw_size} {}
 
  InputData(InputData&& other) noexcept {
    const bool owns_host_buffer = other.ownsHostBuffer();
    input = std::move(other.input);
    data = std::move(other.data);
    data_raw = owns_host_buffer ? data.data() : other.data_raw;
    device_data_raw = other.device_data_raw;
    data_raw_size = other.data_raw_size;
    other.data_raw = nullptr;
    other.device_data_raw = nullptr;
    other.data_raw_size = 0;
  }
 
  InputData& operator=(const InputData& other) {
    if (this != &other) {
      const bool owns_host_buffer = other.ownsHostBuffer();
      input = other.input;
      data = other.data;
      data_raw = owns_host_buffer ? data.data() : other.data_raw;
      device_data_raw = other.device_data_raw;
      data_raw_size = other.data_raw_size;
    }
    return *this;
  }
 
  InputData& operator=(InputData&& other) noexcept {
    if (this != &other) {
      const bool owns_host_buffer = other.ownsHostBuffer();
      input = std::move(other.input);
      data = std::move(other.data);
      data_raw = owns_host_buffer ? data.data() : other.data_raw;
      device_data_raw = other.device_data_raw;
      data_raw_size = other.data_raw_size;
      other.data_raw = nullptr;
      other.device_data_raw = nullptr;
      other.data_raw_size = 0;
    }
    return *this;
  }
 
  InputData(std::shared_ptr<triton::client::InferInput> input, std::vector<uint8_t>&& data)
      : input{std::move(input)},
        data{std::move(data)},
        data_raw{this->data.data()},
        device_data_raw{nullptr},
        data_raw_size{this->data.size()} {}
 
  InputData(std::shared_ptr<triton::client::InferInput> input, uint8_t* data_raw, std::size_t data_raw_size)
      : input{input}, data{}, data_raw{data_raw}, device_data_raw{nullptr}, data_raw_size{data_raw_size} {}
 
  InputData(std::shared_ptr<triton::client::InferInput> input,
            uint8_t* data_raw,
            uint8_t* device_data_raw,
            std::size_t data_raw_size)
      : input{input}, data{}, data_raw{data_raw}, device_data_raw{device_data_raw}, data_raw_size{data_raw_size} {}
 
  bool isHostMappable() const { return data_raw != nullptr; }
  bool isDeviceBacked() const { return device_data_raw != nullptr; }
 
 private:
  bool ownsHostBuffer() const noexcept { return data_raw == data.data(); }
};
 
struct InputOutputMetaData {
  const std::vector<int64_t> shape;
  const inference::DataType datatype;
  const int64_t bytesize;
 
 public:
  InputOutputMetaData(const std::vector<int64_t>& shape, const inference::DataType& datatype)
      : shape{shape},
        datatype{datatype},
        bytesize{accumulate_shape(shape.begin(), shape.end()) *
                 std::visit([](auto&& arg) { return static_cast<std::int64_t>(sizeof(arg)); }, getZero(datatype))} {}
};
 
}  // namespace triton_cpp