stream.h#

namespace ONNX_LIGHT_NAMESPACE

Functions

inline constexpr SerializeSizeResult operator+(SerializeSizeResult left, const SerializeSizeResult &right)#

Returns the sum of two serialized size splits.

Variables

constexpr int64_t kSmallTensorDataThresholdBytes = 8 * static_cast<int64_t>(sizeof(int64_t))#
struct SerializeSizeResult#
#include <stream.h>

Splits serialized bytes between the protobuf payload and separate tensor data.

Public Functions

constexpr SerializeSizeResult() = default#

Initializes an empty size split.

inline constexpr SerializeSizeResult(int64_t small_data_size, int64_t big_data_size, int64_t proto_size)#

Initializes the size split from small data, big data, and protobuf byte counts.

inline constexpr void add_tensor_data_size(int64_t tensor_data_size, int64_t threshold)#

Adds tensor bytes to the small or big bucket depending on threshold.

inline constexpr SerializeSizeResult &operator+=(const SerializeSizeResult &other)#

Accumulates another serialized size split into this result.

inline constexpr int64_t size() const#

Returns the total serialized size across protobuf and external data.

Public Members

int64_t small_data_size = 0#

Stores the number of bytes written to small external tensor data blocks.

int64_t big_data_size = 0#

Stores the number of bytes written to big external tensor data blocks.

int64_t proto_size = 0#

Stores the number of bytes kept in the protobuf payload.

namespace utils

Typedefs

typedef int64_t offset_t#

Signed byte-offset type used by stream seek and length operations.

Functions

inline int64_t decodeZigZag64(uint64_t n)#

Decodes a 64-bit integer from its ZigZag-encoded unsigned representation. ZigZag maps signed integers to unsigned values so that small negative numbers have a short varint encoding: 0 -> 0, -1 -> 1, 1 -> 2, -2 -> 3, …

inline uint64_t encodeZigZag64(int64_t n)#

Encodes a signed 64-bit integer using ZigZag encoding. The resulting unsigned value is suitable for efficient varint serialization of values that are likely to be small in magnitude.

std::shared_ptr<uint8_t> mmap_file_as_shared_ptr(const std::string &file_path, int64_t file_size)#

Maps an entire file into read-only virtual memory and returns a shared_ptr<uint8_t> whose deleter unmaps the region. On POSIX, mmap(MAP_PRIVATE|PROT_READ) is used; on Windows, CreateFileMapping + MapViewOfFile. Returns an empty shared_ptr when file_size is 0. The mapped base address is page-aligned and therefore satisfies any typical tensor alignment requirement (16 / 32 / 64 bytes) when combined with an aligned file offset.

class BinaryStream#
#include <stream.h>

Base class for binary input streams. Concrete subclasses provide byte-level reading from different sources (memory buffers, files) and expose higher-level helpers for varint, float, string, and protobuf tag decoding built on top of the primitives.

Subclassed by ONNX_LIGHT_NAMESPACE::utils::FileStream, ONNX_LIGHT_NAMESPACE::utils::StringStream

Public Functions

inline explicit BinaryStream()#

Initializes a base binary stream with no active limits.

virtual ~BinaryStream()#
inline virtual bool ExternalWeights() const#

Returns true if this stream delivers tensor weights from a separate storage backend.

inline virtual bool CanNoCopy() const#

Returns true if read_bytes(n, nullptr) returns a stable pointer into the backing buffer. When true, the returned pointer remains valid as long as the underlying data buffer lives. In-memory streams (StringStream) return true; file-backed streams return false.

virtual uint64_t next_uint64() = 0#

Reads and decodes the next base-128 varint as an unsigned 64-bit integer.

virtual void CanRead(uint64_t len, const char *msg) = 0#

Raises an exception if fewer than len bytes remain before the current limit. The msg parameter is included in the error message for diagnostics.

virtual bool NotEnd() const = 0#

Returns true if the current read position has not yet reached the active limit.

virtual offset_t tell() const = 0#

Returns the current byte offset from the beginning of the stream.

virtual std::string tell_around() const = 0#

Returns a short human-readable excerpt of bytes around the current position. Useful for error messages and debugging.

virtual const uint8_t *read_bytes(offset_t n_bytes, uint8_t *pre_allocated_buffer = nullptr) = 0#

Reads n_bytes bytes and returns a pointer to the data. If pre_allocated_buffer is non-null the bytes are copied into it; otherwise the stream may return a pointer into its internal buffer.

virtual void skip_bytes(offset_t n_bytes) = 0#

Advances the read position by n_bytes without returning the data.

virtual int64_t size() const = 0#

Returns the total number of bytes in the stream.

virtual RefString next_string()#

Reads the next length-prefixed byte sequence and returns a non-owning view.

virtual int64_t next_int64()#

Reads the next varint and interprets it as a ZigZag-encoded signed 64-bit integer.

virtual int32_t next_int32()#

Reads the next varint and interprets it as a ZigZag-encoded signed 32-bit integer.

virtual float next_float()#

Reads the next 4 bytes and interprets them as a little-endian IEEE 754 float.

virtual double next_double()#

Reads the next 8 bytes and interprets them as a little-endian IEEE 754 double.

virtual FieldNumber next_field()#

Reads the next protobuf tag and returns a FieldNumber struct.

template<typename T>
inline void next_packed_element(T &value)#

Reads a single packed element of type T from the stream into value.

virtual void LimitToNext(uint64_t len)#

Pushes a new read limit of len bytes relative to the current position. NotEnd() will return false once this many bytes have been consumed.

virtual void Restore()#

Pops the most recently pushed read limit, restoring the previous limit.

inline virtual bool HasParallelizationStarted() const#

Returns true once StartThreadPool() has been called and is still active.

virtual void StartThreadPool(size_t n_threads)#

Starts an internal thread pool with n_threads worker threads for parallel block reads.

virtual void ReadDelayedBlock(DelayedBlock &block)#

Submits block to the thread pool so its data is read asynchronously.

virtual void WaitForDelayedBlock()#

Blocks until all pending asynchronous read blocks have completed.

Protected Functions

virtual void LimitTo(uint64_t len) = 0#

Sets the internal read limit to len bytes from the current position.

virtual void _check()#

Validates the internal limit stack; throws if the state is inconsistent.

Protected Attributes

std::vector<uint64_t> limits_#

Stack of absolute byte offsets used to implement nested LimitToNext/Restore pairs.

class BinaryWriteStream#
#include <stream.h>

Base class for binary output streams. Concrete subclasses persist bytes to different backends (memory buffers, files). Higher-level helpers for varint, float, string, and protobuf tag encoding are implemented on top of the write_raw_bytes() primitive.

Subclassed by ONNX_LIGHT_NAMESPACE::utils::BorrowedWriteStream, ONNX_LIGHT_NAMESPACE::utils::FileWriteStream, ONNX_LIGHT_NAMESPACE::utils::StringWriteStream

Public Functions

inline explicit BinaryWriteStream()#

Initializes an empty binary write stream.

inline virtual ~BinaryWriteStream()#
virtual void write_raw_bytes(const uint8_t *data, offset_t n_bytes) = 0#

Appends n_bytes bytes starting at data to the stream.

virtual int64_t size() const = 0#

Returns the number of bytes written so far.

virtual const uint8_t *data() const = 0#

Returns a pointer to the beginning of the stream’s internal buffer. Only valid for in-memory streams; file-backed streams return nullptr.

virtual void write_variant_uint64(uint64_t value)#

Encodes value as a base-128 varint and appends it to the stream.

virtual void write_int64(int64_t value)#

Encodes value as a ZigZag varint and appends it to the stream.

virtual void write_int32(int32_t value)#

Encodes value as a ZigZag varint and appends it to the stream.

virtual void write_float(float value)#

Appends the 4-byte little-endian IEEE 754 representation of value.

virtual void write_double(double value)#

Appends the 8-byte little-endian IEEE 754 representation of value.

virtual void write_string(const std::string &value)#

Encodes value as a length-prefixed byte sequence and appends it.

virtual void write_string(const String &value)#

Encodes value as a length-prefixed byte sequence and appends it.

virtual void write_string(const RefString &value)#

Encodes value as a length-prefixed byte sequence and appends it.

virtual void write_string_stream(const StringWriteStream &stream)#

Encodes the contents of stream as a length-prefixed byte sequence and appends it.

virtual void write_string_stream(const BorrowedWriteStream &stream)#

Encodes the contents of stream as a length-prefixed byte sequence and appends it.

virtual void write_field_header(uint32_t field_number, uint8_t wire_type)#

Encodes a protobuf field tag from field_number and wire_type and appends it.

template<typename T>
inline void write_packed_element(const T &value)#

Appends a single packed element of type T to the stream.

virtual uint64_t size_field_header(uint32_t field_number, uint8_t wire_type)#

Returns the serialized byte size of the protobuf field tag for the given field number and wire type.

virtual uint64_t VarintSize(uint64_t value)#

Returns the number of bytes required to encode value as a base-128 varint.

virtual uint64_t size_variant_uint64(uint64_t value)#

Returns the encoded byte size of value as a varint.

virtual uint64_t size_int64(int64_t value)#

Returns the encoded byte size of value as a ZigZag varint.

virtual uint64_t size_int32(int32_t value)#

Returns the encoded byte size of value as a ZigZag varint.

virtual uint64_t size_float(float value)#

Returns the encoded byte size of value (always 4).

virtual uint64_t size_double(double value)#

Returns the encoded byte size of value (always 8).

virtual uint64_t size_string(const std::string &value)#

Returns the encoded byte size of the length-prefixed string value.

virtual uint64_t size_string(const String &value)#

Returns the encoded byte size of the length-prefixed string value.

virtual uint64_t size_string(const RefString &value)#

Returns the encoded byte size of the length-prefixed string value.

virtual uint64_t size_string_stream(const StringWriteStream &stream)#

Returns the encoded byte size of the contents of stream as a length-prefixed sequence.

virtual uint64_t size_string_stream(const BorrowedWriteStream &stream)#

Returns the encoded byte size of the contents of stream as a length-prefixed sequence.

inline virtual bool ExternalWeights() const#

Returns true if this stream routes tensor weight data to a separate backend.

inline virtual void write_raw_bytes_in_second_stream(const uint8_t*, offset_t)#

Appends n_bytes from data to the external weights backend.

inline virtual void write_raw_bytes_in_second_stream(const uint8_t *data, offset_t n_bytes, const std::string&)#

Appends n_bytes from data to the external weights backend selected by location.

inline virtual int64_t weights_size() const#

Returns number of bytes written to the default external weights backend.

inline virtual int64_t weights_size_for_location(const std::string&) const#

Returns number of bytes written to the external weights backend selected by location.

virtual void CacheSize(const void *ptr, SerializeSizeResult size)#

Associates serialized size information with the object at ptr in the size cache.

virtual bool GetCachedSize(const void *ptr, SerializeSizeResult &size)#

Looks up the cached serialized size for the object at ptr. Returns true and writes the result into size if found.

inline void swap_size_cache(BinaryWriteStream &other)#

Swaps the size cache with other, transferring cached sizes between streams in O(1) so the write pass can reuse sizes computed by a separate size pass.

inline virtual bool HasParallelizationStarted() const#

Returns true once StartThreadPool() has been called and is still active.

virtual void StartThreadPool(size_t n_threads)#

Starts an internal thread pool with n_threads worker threads for parallel block writes.

virtual void WriteDelayedBlock(DelayedWriteBlock &block)#

Submits block to the thread pool so its data is written asynchronously.

virtual void WaitForDelayedBlock()#

Blocks until all pending asynchronous write blocks have completed.

Protected Attributes

std::unordered_map<const void*, SerializeSizeResult> size_cache_#

Per-object serialized-size cache used to avoid redundant recomputation.

class BorrowedStringWriteStream : public ONNX_LIGHT_NAMESPACE::utils::StringWriteStream#
#include <stream.h>

Binary writer backed by a caller-provided fixed-capacity memory buffer. Inherits string-writing helpers from StringWriteStream but never reallocates. Throws std::runtime_error if a write would exceed the initial capacity.

Public Functions

inline explicit BorrowedStringWriteStream(uint8_t *data, int64_t size)#

Initializes a write stream that borrows size bytes starting at data. The caller must ensure the buffer outlives this stream.

virtual void write_raw_bytes(const uint8_t *data, offset_t n_bytes) override#

Appends n_bytes bytes starting at data to the stream.

virtual void WriteDelayedBlock(DelayedWriteBlock &block) override#

Submits block to the thread pool to write directly into the borrowed buffer.

inline virtual int64_t size() const override#

Returns the number of bytes written so far.

inline virtual const uint8_t *data() const override#

Returns a pointer to the beginning of the borrowed buffer.

Protected Attributes

uint8_t *data_#

Non-owning pointer to writable backing storage.

int64_t capacity_#

Maximum number of writable bytes in data_.

class BorrowedWriteStream : public ONNX_LIGHT_NAMESPACE::utils::BinaryWriteStream#
#include <stream.h>

Binary writer backed by externally provided memory. Wraps a caller-owned byte range and exposes the BinaryWriteStream interface without owning or copying the underlying storage.

Public Functions

inline explicit BorrowedWriteStream(const uint8_t *data, int64_t size)#

Initializes a write stream that borrows size bytes starting at data. The caller must ensure the buffer outlives this stream.

virtual void write_raw_bytes(const uint8_t *data, offset_t n_bytes) override#

Appends n_bytes bytes starting at data to the stream.

inline virtual int64_t size() const override#

Returns the number of bytes in the borrowed buffer.

inline virtual const uint8_t *data() const override#

Returns a pointer to the beginning of the borrowed buffer.

Protected Attributes

const uint8_t *data_#

Non-owning pointer to the backing byte buffer.

int64_t size_#

Total number of bytes in the borrowed buffer.

struct DelayedBlock#
#include <stream.h>

Descriptor for a block of data that should be processed asynchronously by a thread pool.

Public Members

uint64_t size#

Number of bytes in the block.

uint8_t *data#

Pointer to the destination buffer that will receive the decoded bytes.

offset_t offset#

Byte offset within the source stream where the block starts.

uint8_t stream_id = 0#

Identifies the substream the data should be read from (0 = primary stream).

struct DelayedWriteBlock#
#include <stream.h>

Descriptor for a block of data that should be written asynchronously by a thread pool.

Public Members

uint64_t size#

Number of bytes to write.

const uint8_t *data#

Pointer to the source data. The caller must keep this buffer alive until WaitForDelayedBlock() returns.

offset_t offset = -1#

Byte offset within the destination stream where the block should be written. A value of -1 means the current sequential write position is used.

uint8_t stream_id = 0#

Identifies the substream the data should be written to (0 = primary stream).

struct FieldNumber#
#include <stream.h>

Decoded protobuf tag containing the field number and wire type.

Public Functions

std::string string() const#

Returns a human-readable representation of the field number and wire type.

Public Members

uint64_t field_number#

Identifies which field in the message this tag belongs to.

uint64_t wire_type#

Wire type that describes the on-wire encoding of the field value.

class FileStream : public ONNX_LIGHT_NAMESPACE::utils::BinaryStream#
#include <stream.h>

Binary reader that streams bytes from a file. Uses a 4096-byte read-ahead buffer (read_buf_) so that sequential varint decoding does not issue a system call per byte. Supports optional parallel block reads via an internal thread pool.

Subclassed by ONNX_LIGHT_NAMESPACE::utils::TwoFilesStream

Public Functions

explicit FileStream(const std::string &file_path)#

Opens the file at file_path for reading.

virtual ~FileStream()#
inline const std::string &file_path() const#

Returns the path of the file this stream reads from.

virtual void CanRead(uint64_t len, const char *msg) override#

Raises an exception if fewer than len bytes remain before the current limit. The msg parameter is included in the error message for diagnostics.

virtual uint64_t next_uint64() override#

Reads and decodes the next base-128 varint as an unsigned 64-bit integer.

virtual const uint8_t *read_bytes(offset_t n_bytes, uint8_t *pre_allocated_buffer = nullptr) override#

Reads n_bytes bytes and returns a pointer to the data. If pre_allocated_buffer is non-null the bytes are copied into it; otherwise the stream may return a pointer into its internal buffer.

virtual void skip_bytes(offset_t n_bytes) override#

Advances the read position by n_bytes without returning the data.

virtual bool NotEnd() const override#

This is a dangerous zone. StreamStream points to the buffer_.data(). buffer_ changes every time new bytes are read from the file. So unlock() must be called or this class raises an exception.

virtual offset_t tell() const override#

Returns the current byte offset from the beginning of the file.

virtual std::string tell_around() const override#

Returns a short human-readable excerpt of bytes around the current position. Useful for error messages and debugging.

virtual bool is_open() const#

Returns true if the underlying file stream is open.

inline virtual int64_t size() const override#

Returns the total number of bytes in the file.

inline virtual bool HasParallelizationStarted() const override#

Returns true once StartThreadPool() has been called and is still active.

virtual void StartThreadPool(size_t n_threads) override#

Starts an internal thread pool with n_threads worker threads for parallel block reads.

virtual void ReadDelayedBlock(DelayedBlock &block) override#

Submits block to the thread pool so its data is read asynchronously.

virtual void WaitForDelayedBlock() override#

Blocks until all pending asynchronous read blocks have completed.

Protected Functions

virtual void LimitTo(uint64_t len) override#

Sets the internal read limit to len bytes from the current position.

void _fill_read_buffer()#

Fills the read-ahead buffer from the current file position.

void _invalidate_read_buffer()#

Seeks the underlying file stream back to tell() and clears the buffer. Must be called before any direct file_stream_ seek or non-buffered read.

Protected Attributes

bool lock_#

Prevents concurrent access while a sub-stream is locked to a region.

std::string file_path_#

Absolute path of the source file.

std::ifstream file_stream_#

Underlying input stream used for file I/O.

int file_descriptor_ = -1#

POSIX file descriptor used for pread-based parallel reads on non-Windows platforms.

int64_t size_#

Total size of the file in bytes.

std::vector<uint8_t> buffer_#

Scratch buffer used to hold bytes read for a single non-buffered read_bytes() call.

std::vector<DelayedBlock> blocks_#

List of blocks scheduled for parallel asynchronous processing.

ThreadPool thread_pool_#

Thread pool used for parallel block reads.

std::vector<uint8_t> read_buf_#

Circular read-ahead buffer of up to READ_BUF_SIZE bytes.

size_t read_buf_pos_ = 0#

Current read index within read_buf_.

size_t read_buf_end_ = 0#

One-past-last valid index within read_buf_.

Protected Static Attributes

static constexpr size_t READ_BUF_SIZE = 4096#

Friends

friend class TwoFilesStream
class FileWriteStream : public ONNX_LIGHT_NAMESPACE::utils::BinaryWriteStream#
#include <stream.h>

Binary writer that persists bytes to a file. Uses an internal std::ofstream with an optional 4096-byte write buffer. Also supports parallel offset-based writes via an internal thread pool after the file has been pre-allocated with pre_allocate().

Subclassed by ONNX_LIGHT_NAMESPACE::utils::TwoFilesWriteStream

Public Functions

explicit FileWriteStream(const std::string &file_path)#

Opens the file at file_path for writing (creates or truncates).

virtual void write_raw_bytes(const uint8_t *data, offset_t n_bytes) override#

Appends n_bytes bytes starting at data to the stream.

virtual int64_t size() const override#

Returns the number of bytes written so far.

virtual const uint8_t *data() const override#

Returns nullptr; file-backed streams do not expose an in-memory buffer.

inline const std::string &file_path() const#

Returns the path of the file this stream writes to.

inline virtual bool HasParallelizationStarted() const override#

Returns true once StartThreadPool() has been called and is still active.

virtual void StartThreadPool(size_t n_threads) override#

Starts an internal thread pool with n_threads worker threads for parallel block writes.

virtual void WriteDelayedBlock(DelayedWriteBlock &block) override#

Submits block to the thread pool so its data is written asynchronously.

virtual void WaitForDelayedBlock() override#

Blocks until all pending asynchronous write blocks have completed.

void pre_allocate(int64_t total_bytes)#

Pre-allocates the file to total_bytes by seeking and writing a zero at the last position. Flushes immediately so the ofstream buffer is clear before parallel tasks write concurrently.

Protected Attributes

std::string file_path_#

Absolute path of the destination file.

std::ofstream file_stream_#

Underlying output stream used for sequential writes.

uint64_t written_bytes_#

Running count of bytes written to the file.

ThreadPool thread_pool_#

Thread pool used for parallel block writes.

class MmapFileStream : public ONNX_LIGHT_NAMESPACE::utils::StringStream#
#include <stream.h>

Binary reader backed by a memory-mapped file. Inherits all in-memory parsing fast paths from StringStream and adds ownership of the mmap region via a shared_ptr so the mapping stays alive as long as this stream (or any borrowed ByteSpan that captured the owner) is in use. Use this in place of FileStream for single-file loads to avoid the double-buffered ifstream-then-read_buf path: the OS page cache is exposed directly as contiguous memory, eliminating per-byte std::ifstream bookkeeping and the seek-to-invalidate sequences triggered by large payload reads.

Public Functions

explicit MmapFileStream(const std::string &file_path)#

Maps the file at file_path and exposes its contents as a binary stream.

inline const std::string &file_path() const#

Returns the path of the mapped file.

inline const std::shared_ptr<uint8_t> &mmap_owner() const#

Returns the shared_ptr that keeps the mmap region alive. Borrowed ByteSpans can capture this pointer so the mapping survives the stream’s destruction, mirroring how TwoFilesStream tracks external weights buffers.

Protected Attributes

std::string file_path_#

Absolute path of the mapped file.

std::shared_ptr<uint8_t> mmap_#

Shared ownership of the mmap region; munmap fires when the last reference dies.

class StringStream : public ONNX_LIGHT_NAMESPACE::utils::BinaryStream#
#include <stream.h>

strings

Binary reader backed by an in-memory buffer. Provides fast sequential access to a caller-owned byte range without copying the underlying data.

Subclassed by ONNX_LIGHT_NAMESPACE::utils::MmapFileStream

Public Functions

inline explicit StringStream()#

Initializes an empty stream pointing to no data.

inline explicit StringStream(const uint8_t *data, int64_t size)#

Initializes a stream that reads size bytes starting at data.

void Setup(const uint8_t *data, int64_t size)#

Resets the stream to read size bytes starting at data.

virtual void CanRead(uint64_t len, const char *msg) override#

Raises an exception if fewer than len bytes remain before the current limit. The msg parameter is included in the error message for diagnostics.

virtual uint64_t next_uint64() override#

Reads and decodes the next base-128 varint as an unsigned 64-bit integer.

virtual const uint8_t *read_bytes(offset_t n_bytes, uint8_t *pre_allocated_buffer = nullptr) override#

Reads n_bytes bytes and returns a pointer to the data. If pre_allocated_buffer is non-null the bytes are copied into it; otherwise the stream may return a pointer into its internal buffer.

virtual void skip_bytes(offset_t n_bytes) override#

Advances the read position by n_bytes without returning the data.

inline virtual bool NotEnd() const override#

Returns true while the read position is before the end of the active limit.

inline virtual bool CanNoCopy() const override#

Returns true: read_bytes(n, nullptr) yields a stable pointer into data_.

inline virtual offset_t tell() const override#

Returns the current byte offset from the beginning of the buffer.

virtual std::string tell_around() const override#

Returns a short human-readable excerpt of bytes around the current position. Useful for error messages and debugging.

inline virtual int64_t size() const override#

Returns the total number of bytes in the buffer.

inline virtual bool HasParallelizationStarted() const override#

Returns true once StartThreadPool() has been called and is still active.

virtual void StartThreadPool(size_t n_threads) override#

Starts an internal thread pool with n_threads worker threads for parallel block reads.

virtual void ReadDelayedBlock(DelayedBlock &block) override#

Submits block to the thread pool so its data is read asynchronously.

virtual void WaitForDelayedBlock() override#

Blocks until all pending asynchronous read blocks have completed.

Protected Functions

virtual void LimitTo(uint64_t len) override#

Sets the internal read limit to len bytes from the current position.

Protected Attributes

offset_t pos_#

Current read offset within data_.

offset_t size_#

Active read limit; equals the buffer size when no sub-limit is active.

const uint8_t *data_#

Non-owning pointer to the backing byte buffer.

std::vector<DelayedBlock> blocks_#

List of blocks scheduled for parallel asynchronous processing.

ThreadPool thread_pool_#

Thread pool used for parallel block reads.

Friends

friend class FileStream
class StringWriteStream : public ONNX_LIGHT_NAMESPACE::utils::BinaryWriteStream#
#include <stream.h>

Binary writer backed by an owned memory buffer. All bytes are accumulated in an internal std::vector<uint8_t>. Supports optional parallel writes via an internal thread pool after the buffer has been pre-allocated with pre_allocate().

Subclassed by ONNX_LIGHT_NAMESPACE::utils::BorrowedStringWriteStream

Public Functions

inline explicit StringWriteStream()#

Initializes an empty write stream with no pre-allocated storage.

virtual void write_raw_bytes(const uint8_t *data, offset_t n_bytes) override#

Appends n_bytes bytes starting at data to the stream.

virtual int64_t size() const override#

Returns the number of bytes written so far.

virtual const uint8_t *data() const override#

Returns a pointer to the beginning of the internal buffer.

void pre_allocate(int64_t total_bytes)#

Pre-allocates the buffer to total_bytes bytes (zero-filled). Requires calling before StartThreadPool; ensures buffer_.data() remains stable so no reallocation occurs during concurrent writes.

inline virtual bool HasParallelizationStarted() const override#

Returns true once StartThreadPool() has been called and is still active.

virtual void StartThreadPool(size_t n_threads) override#

Starts an internal thread pool with n_threads worker threads for parallel block writes.

virtual void WriteDelayedBlock(DelayedWriteBlock &block) override#

Submits block to the thread pool so its data is written asynchronously.

virtual void WaitForDelayedBlock() override#

Blocks until all pending asynchronous write blocks have completed.

Protected Attributes

std::vector<uint8_t> buffer_#

Owned byte buffer that accumulates written data.

offset_t write_pos_#

Current sequential write offset within buffer_.

ThreadPool thread_pool_#

Thread pool used for parallel block writes.

class TwoFilesStream : public ONNX_LIGHT_NAMESPACE::utils::FileStream#
#include <stream.h>

Two-file reader for ONNX models with external tensor data. Reads the model protobuf from a primary file (inherited from FileStream) and tensor weight data from a separate weights file on demand.

Public Functions

explicit TwoFilesStream(const std::string &file_path, const std::string &weights_file)#

Opens file_path for protobuf data and weights_file for weight data.

inline const std::string &weights_file_path() const#

Returns the path of the separate weights file.

void set_active_weights_location(const std::string &location)#

Selects the active external weights location for subsequent reads.

bool using_default_weights_location() const#

Returns true when the active location is the default weights file.

inline uint64_t weights_tell() const#

Returns the current byte offset within the weights file.

inline virtual bool ExternalWeights() const override#

Returns true; this stream reads tensor data from a separate weights file.

virtual void read_bytes_from_weights_stream(offset_t n_bytes, uint8_t *pre_allocated_buffer = nullptr, offset_t offset = -1)#

Reads n_bytes from the weights file at the given offset (or sequentially if -1) into pre_allocated_buffer.

const uint8_t *borrow_weights_bytes(const std::string &location, offset_t offset, offset_t n_bytes, size_t alignment, std::shared_ptr<void> &owner)#

Loads one external weights file once, then returns a shared borrowed view into it.

virtual void ReadDelayedBlock(DelayedBlock &block) override#

Submits block to the thread pool so its data is read asynchronously.

inline virtual int64_t weights_size() const#

Returns the total number of bytes in the weights file.

int64_t weights_size(const std::string &location) const#

Returns the total number of bytes in one weights file.

Protected Functions

FileStream &active_weights_stream()#
const FileStream &active_weights_stream() const#
SharedWeightsBuffer &ensure_shared_weights_buffer(const std::string &location, size_t alignment)#

Protected Attributes

FileStream weights_stream_#

Reader for the separate weights file.

std::unordered_map<std::string, std::unique_ptr<FileStream>> extra_weights_streams_#

Additional readers when external_data.location points to multiple files.

std::unordered_map<std::string, SharedWeightsBuffer> shared_weights_buffers_#

Shared file buffers used by no-copy external-data loading.

std::string active_weights_location_#

Active external location used by the current tensor read.

std::string default_weights_location_#

Relative location key associated with the default weights stream.

std::unordered_map<const void*, uint64_t> position_cache_#

Maps object pointers to their byte offsets in the weights file.

struct SharedWeightsBuffer#
#include <stream.h>

Public Members

std::shared_ptr<uint8_t> data#
int64_t size = 0#
size_t alignment = 0#
class TwoFilesWriteStream : public ONNX_LIGHT_NAMESPACE::utils::FileWriteStream#
#include <stream.h>

Two-file writer for external ONNX tensor data. Buffers the model protobuf structure in memory (main_buf_) and flushes it to the primary file in a single write via FlushMainToFile(). Tensor weight data is streamed directly to a separate weights file. This separates the two I/O streams so the OS can coalesce writes efficiently and avoids the overhead of many small writes to the main file. Supports parallel offset-based writes to the weights file via StartWriteThreadPool().

Public Functions

explicit TwoFilesWriteStream(const std::string &file_path, const std::string &weights_file)#

Opens file_path for protobuf data and weights_file for weight data.

inline const std::string &weights_file_path() const#

Returns the path of the separate weights file.

void set_active_weights_location(const std::string &location)#

Selects the active external weights location for subsequent tensor raw-data writes.

inline virtual bool ExternalWeights() const override#

Returns true; this stream routes tensor data to a separate weights file.

virtual void write_raw_bytes_in_second_stream(const uint8_t *data, offset_t n_bytes) override#

Appends n_bytes bytes starting at data to the weights file.

virtual void write_raw_bytes_in_second_stream(const uint8_t *data, offset_t n_bytes, const std::string &location) override#

Appends n_bytes bytes to the weights file designated by location.

virtual int64_t weights_size() const override#

Returns the number of bytes written to the weights file so far.

int64_t weights_size(const std::string &location) const#

Returns the number of bytes written so far for one weights location.

virtual int64_t weights_size_for_location(const std::string &location) const override#

Returns the number of bytes written so far for one weights location.

void pre_allocate_weights(int64_t total_bytes)#

Pre-allocates the weights file to total_bytes by writing a zero at the last position. Must be called before StartWriteThreadPool.

void StartWriteThreadPool(int32_t n_threads)#

Starts a thread pool of n_threads workers for parallel offset-based writes. After this call write_raw_bytes_in_second_stream submits writes asynchronously.

void WaitForWriteCompletion()#

Blocks until all pending write tasks have completed and stops the thread pool.

void FlushMainToFile()#

Flushes the in-memory main-file buffer to disk in a single write. Must be called after all serialization is complete (after WaitForWriteCompletion).

virtual void write_raw_bytes(const uint8_t *data, offset_t n_bytes) override#

Writes n_bytes bytes from data into the in-memory main-content buffer. Overrides the base file write so that the two output files are kept separate and can be flushed independently.

virtual int64_t size() const override#

Returns the number of bytes accumulated in the main-content buffer so far.

inline virtual void StartThreadPool(size_t) override#

Does nothing: the main-content buffer grows dynamically and does not need thread-pool parallelism because all large tensor writes go to the weights stream.

inline virtual bool HasParallelizationStarted() const override#

Returns false because no parallel writes are submitted for the main-content buffer.

virtual void WriteDelayedBlock(DelayedWriteBlock &block) override#

Throws unconditionally: main-content writes are sequential and this path is unreachable.

inline virtual void WaitForDelayedBlock() override#

Does nothing because no delayed writes are outstanding for the main-content buffer.

Protected Attributes

StringWriteStream main_buf_#

In-memory buffer that accumulates the main .onnx structure bytes. Flushed to the primary file in one shot by FlushMainToFile().

FileWriteStream weights_stream_#

Writer for the separate weights file.

std::unordered_map<std::string, std::unique_ptr<FileWriteStream>> extra_weights_streams_#

Additional writers when external_data.location points to multiple files.

std::string active_weights_location_#

Active external location used by the current tensor raw-data write.

std::string default_weights_location_#

Relative location key associated with the default weights stream.

bool parallel_write_ = false#

Set to true once StartWriteThreadPool() has been called.

int64_t virtual_write_pos_ = 0#

Tracks the sequential write position for offset validation during parallel writes.

ThreadPool write_thread_pool_#

Thread pool used for parallel writes to the weights file.