From 59db4617e1cf8d1f8bbe9aecabc52f1ffec95557 Mon Sep 17 00:00:00 2001 From: Theresa Pollinger Date: Thu, 19 Feb 2026 17:30:02 +0900 Subject: [PATCH] TensorFile: add loading from memory --- src/dalotia.cpp | 26 +++++++++++++++++++++++++- src/dalotia.f90 | 8 ++++++++ src/dalotia.h | 2 ++ src/dalotia.hpp | 2 ++ src/dalotia_safetensors_file.cpp | 21 +++++++++++++++++++++ src/dalotia_safetensors_file.hpp | 2 ++ 6 files changed, 60 insertions(+), 1 deletion(-) diff --git a/src/dalotia.cpp b/src/dalotia.cpp index 3f7946a..690f2be 100644 --- a/src/dalotia.cpp +++ b/src/dalotia.cpp @@ -48,7 +48,7 @@ TensorFile *make_tensor_file(const std::string &filename) { ::tolower); // select the file implementation - if (extension == "safetensors") { + if (extension == "safetensors" || extension == "st") { #ifdef DALOTIA_WITH_SAFETENSORS_CPP return new SafetensorsFile(filename); #else // DALOTIA_WITH_SAFETENSORS_CPP @@ -66,6 +66,24 @@ TensorFile *make_tensor_file(const std::string &filename) { return nullptr; } + +// factory function for the file, selected by file extension and +// available implementations +TensorFile *load_tensor_file_from_memory(const void * const address, size_t num_bytes, const std::string &format) { + auto& extension = format; + // select the file implementation + if (extension == "safetensors") { +#ifdef DALOTIA_WITH_SAFETENSORS_CPP + return new SafetensorsFile(address, num_bytes); +#else // DALOTIA_WITH_SAFETENSORS_CPP + throw std::runtime_error("Safetensors support not enabled"); +#endif // DALOTIA_WITH_SAFETENSORS_CPP + } else { + throw std::runtime_error("Unsupported memory format: ." + extension); + } + return nullptr; +} + } // namespace dalotia DalotiaTensorFile *dalotia_open_file(const char *filename) { @@ -73,6 +91,12 @@ DalotiaTensorFile *dalotia_open_file(const char *filename) { dalotia::make_tensor_file(std::string(filename))); } +DalotiaTensorFile *dalotia_load_file_from_memory(const void * const address, size_t num_bytes, const char *format) { + return reinterpret_cast( + dalotia::load_tensor_file_from_memory(address, num_bytes, std::string(format))); +} + + void dalotia_close_file(DalotiaTensorFile *file) { delete reinterpret_cast(file); } diff --git a/src/dalotia.f90 b/src/dalotia.f90 index e21e564..9d19963 100644 --- a/src/dalotia.f90 +++ b/src/dalotia.f90 @@ -33,6 +33,14 @@ type(C_ptr) function dalotia_open_file_c(file_name) bind(C,name="dalotia_open_fi character(kind=C_char), dimension(*), intent(in):: file_name end function dalotia_open_file_c + type(C_ptr) function dalotia_load_file_from_memory_c(address, num_bytes, file_format) bind(C,name="dalotia_load_file_from_memory") + use, intrinsic::ISO_C_BINDING, only: C_ptr, C_char, C_size_t + implicit none + type(C_ptr), intent(in), value :: address + integer(C_size_t) :: num_bytes + character(kind=C_char), dimension(*), intent(in):: file_format + end function dalotia_load_file_from_memory_c + subroutine dalotia_close_file(dalotia_file_pointer) bind(C,name="dalotia_close_file") use, intrinsic::ISO_C_BINDING, only: C_ptr implicit none diff --git a/src/dalotia.h b/src/dalotia.h index 749c491..b50f859 100644 --- a/src/dalotia.h +++ b/src/dalotia.h @@ -17,6 +17,8 @@ typedef struct DalotiaTensorFile DalotiaTensorFile; EXTERNC DalotiaTensorFile *dalotia_open_file(const char *filename); +EXTERNC DalotiaTensorFile *dalotia_load_file_from_memory(const void *address, size_t num_bytes, const char *format); + EXTERNC void dalotia_close_file(DalotiaTensorFile *file); EXTERNC int dalotia_sizeof_weight_format(dalotia_WeightFormat format); diff --git a/src/dalotia.hpp b/src/dalotia.hpp index 8b96cc7..aaa2008 100644 --- a/src/dalotia.hpp +++ b/src/dalotia.hpp @@ -26,6 +26,8 @@ namespace dalotia { // available implementations [[nodiscard]] TensorFile *make_tensor_file(const std::string & filename); +[[nodiscard]] TensorFile *load_tensor_file_from_memory(const void * const address, size_t num_bytes, const char *format); + // C++17 version -> will not compile on Fugaku... // -- pmr vector types can accept different allocators //? more memory interface than that? detect if CUDA device pointer through diff --git a/src/dalotia_safetensors_file.cpp b/src/dalotia_safetensors_file.cpp index e7dc267..a120e84 100644 --- a/src/dalotia_safetensors_file.cpp +++ b/src/dalotia_safetensors_file.cpp @@ -71,6 +71,27 @@ SafetensorsFile::SafetensorsFile(const std::string &filename) : TensorFile(filen #endif // NDEBUG } +SafetensorsFile::SafetensorsFile(const void * const address, size_t num_bytes) : TensorFile("") { + // as far as I can tell, safetensors are saved in C order + std::string warn, err; + bool ret = safetensors::mmap_from_memory(static_cast(address), num_bytes, "", &st_, &warn, &err); + if (warn.size() > 0) { + std::cout << "safetensors-cpp WARN: " << warn << "\n"; + } + if (ret == false) { + std::cerr << " ERR: " << err << "\n"; + throw std::runtime_error("Could not load safetensors from address"); + } +#ifndef NDEBUG + // Check if data_offsets are valid + if (!safetensors::validate_data_offsets(st_, err)) { + std::cerr << "Invalid data_offsets\n"; + std::cerr << err << "\n"; + throw std::runtime_error("Invalid safetensors address"); + } +#endif // NDEBUG +} + SafetensorsFile::~SafetensorsFile() { if (st_.st_file != nullptr) { // delete st_.st_file; diff --git a/src/dalotia_safetensors_file.hpp b/src/dalotia_safetensors_file.hpp index b8fd0d6..0f88546 100644 --- a/src/dalotia_safetensors_file.hpp +++ b/src/dalotia_safetensors_file.hpp @@ -31,6 +31,8 @@ class SafetensorsFile : public TensorFile { public: explicit SafetensorsFile(const std::string &filename); + SafetensorsFile(const void * const address, size_t num_bytes); + ~SafetensorsFile() override; const std::vector &get_tensor_names() const override;