Implementation:LaurentMazare Tch rs Stream IO

Knowledge Sources	LaurentMazare_Tch_rs
Domains	FFI, Serialization, I/O
Last Updated	2026-02-08 00:00 GMT

Overview

The Stream I/O module provides extern "C" callback functions that bridge Rust's Read, Write, and Seek traits to the C++ LibTorch serialization layer, enabling tensor save/load through arbitrary Rust I/O streams.

Description

This module in torch-sys defines a set of #[no_mangle] extern "C" functions that serve as C-callable callbacks for the LibTorch stream-based serialization API. The C++ side holds an opaque *mut c_void pointer to a double-boxed Rust trait object (Box<Box<dyn Write>> or Box<Box<dyn ReadStream>>), and calls these functions to perform I/O operations.

The ReadStream trait is a convenience alias combining Read + Seek, required because LibTorch's deserialization needs both sequential reading and random seeking capabilities.

Each callback follows a consistent pattern: cast the opaque pointer back to the trait object, invoke the corresponding Rust trait method, write the result through an output pointer, and return a bool indicating success or failure. Destructors are provided to properly drop the boxed trait objects when the C++ side is done with them.

The six callback functions are:

• tch_write_stream_destructor - Drops the boxed dyn Write trait object
• tch_write_stream_write - Writes a buffer through the Write trait
• tch_read_stream_destructor - Drops the boxed dyn ReadStream trait object
• tch_read_stream_stream_position - Queries the current stream position via Seek::stream_position
• tch_read_stream_seek_start - Seeks to an absolute position via SeekFrom::Start
• tch_read_stream_seek_end - Seeks relative to the end via SeekFrom::End
• tch_read_stream_read - Reads bytes through the Read trait

Usage

These callbacks are not called directly by user code. They are registered with the C++ LibTorch runtime when saving or loading tensors through Rust stream objects (e.g., File, BufWriter, or in-memory Cursor). The higher-level tch crate wraps these into safe APIs for tensor serialization.

Code Reference

Source Location

• Repository: LaurentMazare_Tch_rs
• File: torch-sys/src/io.rs

Signature

pub trait ReadStream: Read + Seek {}

extern "C" fn tch_write_stream_destructor(stream_ptr: *mut c_void);

extern "C" fn tch_write_stream_write(
    stream_ptr: *mut c_void,
    buf: *const u8,
    size: size_t,
    out_size: *mut size_t,
) -> bool;

extern "C" fn tch_read_stream_destructor(stream_ptr: *mut c_void);

extern "C" fn tch_read_stream_stream_position(
    stream_ptr: *mut c_void,
    current: *mut u64,
) -> bool;

extern "C" fn tch_read_stream_seek_start(
    stream_ptr: *mut c_void,
    pos: u64,
    new_pos: *mut u64,
) -> bool;

extern "C" fn tch_read_stream_seek_end(
    stream_ptr: *mut c_void,
    pos: i64,
    new_pos: *mut u64,
) -> bool;

extern "C" fn tch_read_stream_read(
    stream_ptr: *mut c_void,
    buf: *mut u8,
    size: size_t,
    read_size: *mut size_t,
) -> bool;

Import

use torch_sys::io::ReadStream;

I/O Contract

Write Callback

Parameter	Type	Description
stream_ptr	*mut c_void	Opaque pointer to Box<dyn Write>
buf	*const u8	Source buffer to write from
size	size_t	Number of bytes in the buffer
out_size	*mut size_t	Output: number of bytes actually written
Return	Type	Description
Success	true	Write succeeded, out_size is populated
Failure	false	Write returned an I/O error

Read Callback

Parameter	Type	Description
stream_ptr	*mut c_void	Opaque pointer to Box<dyn ReadStream>
buf	*mut u8	Destination buffer to read into
size	size_t	Maximum number of bytes to read
read_size	*mut size_t	Output: number of bytes actually read
Return	Type	Description
Success	true	Read succeeded, read_size is populated
Failure	false	Read returned an I/O error

Seek Callbacks

Function	Seek Mode	Position Parameter	Output
tch_read_stream_stream_position	Current position	N/A	current: *mut u64
tch_read_stream_seek_start	SeekFrom::Start	pos: u64 (absolute)	new_pos: *mut u64
tch_read_stream_seek_end	SeekFrom::End	pos: i64 (relative to end)	new_pos: *mut u64

All seek callbacks return bool: true on success with the output pointer populated, false on I/O error.

Destructor Callbacks

Function	Drops
tch_write_stream_destructor	Box<Box<dyn Write>>
tch_read_stream_destructor	Box<Box<dyn ReadStream>>

Usage Examples

// The callbacks are not called directly. They are used internally
// when saving/loading tensors through Rust streams. For example,
// at the higher-level tch crate:

use tch::Tensor;
use std::io::Cursor;

// Save a tensor to an in-memory buffer (internally uses the write callbacks)
let tensor = Tensor::randn(&[3, 4], (tch::Kind::Float, tch::Device::Cpu));
let mut buffer = Vec::new();
tensor.save_to_stream(&mut buffer).unwrap();

// Load a tensor from an in-memory buffer (internally uses the read callbacks)
let mut cursor = Cursor::new(buffer);
let loaded = Tensor::load_from_stream(cursor).unwrap();

// Implementing ReadStream for a custom type:
use torch_sys::io::ReadStream;
use std::io::{Read, Seek};

struct MyStream { /* ... */ }

impl Read for MyStream {
    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> { /* ... */ }
}

impl Seek for MyStream {
    fn seek(&mut self, pos: std::io::SeekFrom) -> std::io::Result<u64> { /* ... */ }
}

impl ReadStream for MyStream {}

Related Pages

• Principle:LaurentMazare_Tch_rs_Stream_Serialization