metatensor/io/
mod.rs

1//! Input/Output facilities for storing [`crate::TensorMap`] and
2//! [`crate::Labels`] on disk
3
4use std::os::raw::c_void;
5
6use dlpk::sys::{DLDataType, DLDataTypeCode};
7
8use crate::c_api::{MTS_SUCCESS, mts_array_t, mts_status_t};
9use crate::MtsArray;
10
11mod tensor;
12pub use self::tensor::{load, load_custom_array, load_buffer, load_buffer_custom_array};
13pub use self::tensor::{save, save_buffer};
14
15mod block;
16pub use self::block::{load_block, load_block_buffer_custom_array, load_block_buffer, load_block_custom_array};
17pub use self::block::{save_block, save_block_buffer};
18
19mod labels;
20pub use self::labels::{load_labels, load_labels_buffer, save_labels, save_labels_buffer};
21
22
23/// Implementation of realloc for `Vec<u8>`, used in `save_buffer`
24unsafe extern "C" fn realloc_vec(user_data: *mut c_void, _ptr: *mut u8, new_size: usize) -> *mut u8 {
25    let mut result = std::ptr::null_mut();
26    let unwind_wrapper = std::panic::AssertUnwindSafe(&mut result);
27
28    let status = crate::errors::catch_unwind(move || {
29        let vector = unsafe { &mut *user_data.cast::<Vec<u8>>() };
30        vector.resize(new_size, 0);
31
32        // force the closure to capture the full unwind_wrapper, not just
33        // unwind_wrapper.0
34        let _ = &unwind_wrapper;
35        *(unwind_wrapper.0) = vector.as_mut_ptr();
36
37        Ok(())
38    });
39
40    if status != MTS_SUCCESS {
41        return std::ptr::null_mut();
42    }
43
44    return result;
45}
46
47/// Create a typed `ndarray::Array<T>` and put it in a `MtsArray`
48macro_rules! create_typed_array {
49    ($shape:expr, $c_array:expr, $T:ty) => {{
50        let array = ndarray::Array::<$T, _>::from_elem($shape, <$T>::default());
51        std::convert::Into::<MtsArray>::into(array)
52    }};
53}
54
55/// Callback function to create a new `mts_array_t` when loading `TensorMap` or
56/// `TensorBlock` from a buffer or a file.
57///
58/// This is an implementation of `mts_create_array_callback_t` that creates a
59/// new `ndarray::ArrayD` as the backing array for the `mts_array_t`. It is used
60/// by default in [`load_buffer`] and [`load_block_buffer`].
61pub unsafe extern "C" fn create_ndarray(shape: *const usize, shape_count: usize, dtype: DLDataType, array: *mut mts_array_t) -> mts_status_t {
62    if dtype.lanes != 1 {
63        let error = crate::Error {
64            code: None,
65            message: format!(
66                "unsupported dtype in create_ndarray: lanes={} (expected 1)",
67                dtype.lanes
68            ),
69        };
70        return crate::errors::store_last_error(error);
71    }
72
73    let shape = unsafe {
74        std::slice::from_raw_parts(shape, shape_count)
75    };
76
77    let new_array = match (dtype.code, dtype.bits) {
78        (DLDataTypeCode::kDLFloat, 32) => create_typed_array!(shape, c_array, f32),
79        (DLDataTypeCode::kDLFloat, 64) => create_typed_array!(shape, c_array, f64),
80        (DLDataTypeCode::kDLInt, 8) => create_typed_array!(shape, c_array, i8),
81        (DLDataTypeCode::kDLInt, 16) => create_typed_array!(shape, c_array, i16),
82        (DLDataTypeCode::kDLInt, 32) => create_typed_array!(shape, c_array, i32),
83        (DLDataTypeCode::kDLInt, 64) => create_typed_array!(shape, c_array, i64),
84        (DLDataTypeCode::kDLUInt, 8) => create_typed_array!(shape, c_array, u8),
85        (DLDataTypeCode::kDLUInt, 16) => create_typed_array!(shape, c_array, u16),
86        (DLDataTypeCode::kDLUInt, 32) => create_typed_array!(shape, c_array, u32),
87        (DLDataTypeCode::kDLUInt, 64) => create_typed_array!(shape, c_array, u64),
88        (DLDataTypeCode::kDLBool, 8) => create_typed_array!(shape, c_array, bool),
89        (DLDataTypeCode::kDLFloat, 16) => create_typed_array!(shape, c_array, half::f16),
90        (DLDataTypeCode::kDLComplex, 64) => create_typed_array!(shape, c_array, [f32; 2]),
91        (DLDataTypeCode::kDLComplex, 128) => create_typed_array!(shape, c_array, [f64; 2]),
92        _ => {
93            let error = crate::Error {
94                code: None,
95                message: format!(
96                    "unsupported dtype in create_ndarray: code={:?} bits={}",
97                    dtype.code, dtype.bits
98                ),
99            };
100            return crate::errors::store_last_error(error);
101        }
102    };
103
104    unsafe {
105        *array = new_array.into_raw();
106    }
107
108    return MTS_SUCCESS;
109}
metatensor/io/mod.rs

metatensor/io/
mod.rs
