vuron/adept: Платформа для обучения и вывода нейронных сетей, с API приближенным к PyTorch.

src/serialize/fileinput.cpp
146 строк4 KB
kolkir
revert compiler compatibility to gcc12
03 мар 2025, 23:52
03 мар 2025, 23:5226a9274
#include <adept/serialize/fileinput.hpp>
#include <adept/tensor.hpp>
#include <adept/types_dispatch.hpp>

#include <cstring>
#include <type_traits>

#include <fcntl.h>
#include <sys/mman.h>
#include <sys/stat.h>

namespace adept {

namespace {

template <typename T>
requires std::is_integral_v<T> || std::is_floating_point_v<T> size_t read_number(T& value,
                                                                                 char* ptr) {
  memcpy(&value, ptr, sizeof(T));
  return sizeof(T);
}

size_t read_string(std::string& str, char* ptr) {
  size_t len = 0;
  ptr += read_number(len, ptr);
  str.resize(len);
  memcpy(str.data(), ptr, len);
  return sizeof(size_t) + len;
}

// size_t read_buffer(std::vector<std::byte>& buf, std::istream& os) {
//   size_t len = 0;
//   read_len(len, os);
//   buf.resize(len);
//   os.read(reinterpret_cast<char*>(buf.data()), len);
// }

}  // namespace

FileInput::FileInput(const std::string& file_name) {
  fd_ = open(file_name.c_str(), O_RDONLY);
  if (fd_ < 0)
    THROW_ERROR("Failed to open file {}", file_name);

  struct stat file_info;
  memset(&file_info, 0, sizeof(file_info));

  if (fstat(fd_, &file_info) == -1) {
    close(fd_);
    THROW_ERROR("Failed to get file info for {}", file_name);
  }
  data_size_ = file_info.st_size;

  data_ptr_ = reinterpret_cast<char*>(mmap(0, data_size_, PROT_READ, MAP_SHARED, fd_, 0));
  if (data_ptr_ == MAP_FAILED) {
    close(fd_);
    THROW_ERROR("Failed to mmap file {}", file_name);
  }

  // read all blocks and remember their offsets
  data_map_.clear();
  char* current_ptr = data_ptr_;
  char* last_ptr = data_ptr_ + data_size_;
  while (current_ptr < last_ptr) {
    auto inner_ptr = current_ptr;
    size_t block_size = 0;
    inner_ptr += read_number(block_size, inner_ptr);
    std::string name;
    inner_ptr += read_string(name, inner_ptr);
    data_map_.emplace(std::move(name), inner_ptr);
    // go to next block
    current_ptr += block_size;
  }
}

FileInput::~FileInput() {
  if (data_ptr_) {
    if (munmap(data_ptr_, data_size_) < 0) {
      std::cerr << "FileInput failed to unmap input file" << std::endl;
    }
  }
  if (fd_ >= 0) {
    if (close(fd_) < 0) {
      std::cerr << "FileInput failed to close input file" << std::endl;
    }
  }
}

void FileInput::read(const std::string& name, float32_t& value) {
  auto i = data_map_.find(name);
  if (i == data_map_.end())
    THROW_ERROR("There is no field ", name, " in FileInput");
  read_number(value, i->second);
}

void FileInput::read(const std::string& name, float64_t& value) {
  auto i = data_map_.find(name);
  if (i == data_map_.end())
    THROW_ERROR("There is no field ", name, " in FileInput");
  read_number(value, i->second);
}

void FileInput::read(const std::string& name, index_t& value) {
  auto i = data_map_.find(name);
  if (i == data_map_.end())
    THROW_ERROR("There is no field ", name, " in FileInput");
  read_number(value, i->second);
}

void FileInput::read(const std::string& name, int32_t& value) {
  auto i = data_map_.find(name);
  if (i == data_map_.end())
    THROW_ERROR("There is no field ", name, " in FileInput");
  read_number(value, i->second);
}

void FileInput::read(const std::string& name, int8_t& value) {
  auto i = data_map_.find(name);
  if (i == data_map_.end())
    THROW_ERROR("There is no field ", name, " in FileInput");
  read_number(value, i->second);
}

void FileInput::read(const std::string& name, Tensor& tensor) {
  auto i = data_map_.find(name);
  if (i == data_map_.end())
    THROW_ERROR("There is no field ", name, " in FileInput");

  auto buffer_ptr = i->second;

  DISPATCH_TYPE(tensor.properties().dtype, [&]() {
    auto numel = tensor.properties().shape.numel();
    size_t buffer_size = 0;
    buffer_ptr += read_number(buffer_size, buffer_ptr);
    if (numel * sizeof(scalar_t) != buffer_size)
      THROW_ERROR("Tensor size is different for ", name, " in FileInput");

    memcpy(tensor.mutable_data_ptr<scalar_t>(), buffer_ptr, buffer_size);
  });
}

void FileInput::read(const std::string& name, Variable& variable) {
  read(name, variable.data());
};

}  // namespace adept