tenferro_runtime/runtime/
signature.rs

1use std::mem::align_of;
2use std::mem::size_of;
3
4use tenferro_tensor::{
5    AllocationDomainId, DType, Placement, ShapeVec, StrideVec, Tensor, TensorRead, TensorScalar,
6    TensorView, TypedTensor, TypedTensorView,
7};
8
9use super::{InputSignatureError, LayoutClass, PrepareError};
10
11const COMPACT_COL_MAJOR_LAYOUT: &str = "tenferro.layout.compact-col-major.v1";
12const STRIDED_LAYOUT: &str = "tenferro.layout.strided.v1";
13
14/// Value-free metadata signature for a tensor input.
15///
16/// # Examples
17///
18/// ```
19/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
20/// use tenferro_runtime::{DType, InputSignatureEntry, LayoutClass};
21/// use tenferro_tensor::Placement;
22///
23/// let entry = InputSignatureEntry::new(
24///     DType::F64,
25///     [2_usize].into_iter().collect(),
26///     Placement::default(),
27///     LayoutClass::new("tenferro.layout.strided")?,
28///     [1_isize].into_iter().collect(),
29///     Some(3),
30/// )?;
31/// assert_eq!(entry.dtype(), DType::F64);
32/// # Ok(())
33/// # }
34/// ```
35#[derive(Clone, Debug, Eq, Hash, PartialEq)]
36pub struct InputSignatureEntry {
37    dtype: DType,
38    shape: ShapeVec,
39    placement: Placement,
40    layout_class: LayoutClass,
41    strides: StrideVec,
42    alignment_log2: Option<u8>,
43    backend_family: Option<&'static str>,
44    allocation_domain: Option<AllocationDomainId>,
45}
46
47#[derive(Clone, Copy)]
48struct InputPhysicalIdentity {
49    backend_family: Option<&'static str>,
50    allocation_domain: Option<AllocationDomainId>,
51}
52
53impl InputSignatureEntry {
54    /// Build one value-free input signature entry.
55    ///
56    /// # Examples
57    ///
58    /// ```
59    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
60    /// use tenferro_runtime::{DType, InputSignatureEntry, LayoutClass};
61    /// use tenferro_tensor::Placement;
62    ///
63    /// let entry = InputSignatureEntry::new(
64    ///     DType::I32,
65    ///     [4_usize].into_iter().collect(),
66    ///     Placement::default(),
67    ///     LayoutClass::new("tenferro.layout.compact")?,
68    ///     [1_isize].into_iter().collect(),
69    ///     None,
70    /// )?;
71    /// assert_eq!(entry.shape(), &[4]);
72    /// # Ok(())
73    /// # }
74    /// ```
75    ///
76    /// # Errors
77    ///
78    /// Returns [`InputSignatureError::ShapeStrideRankMismatch`] when shape and
79    /// stride ranks differ, or [`InputSignatureError::InvalidAlignmentClass`]
80    /// when `alignment_log2` is outside the finite `usize` alignment lattice.
81    pub fn new(
82        dtype: DType,
83        shape: ShapeVec,
84        placement: Placement,
85        layout_class: LayoutClass,
86        strides: StrideVec,
87        alignment_log2: Option<u8>,
88    ) -> Result<Self, InputSignatureError> {
89        validate_entry(&shape, &strides, alignment_log2)?;
90        Ok(Self {
91            dtype,
92            shape,
93            placement,
94            layout_class,
95            strides,
96            alignment_log2,
97            backend_family: None,
98            allocation_domain: None,
99        })
100    }
101
102    fn from_validated_metadata(
103        dtype: DType,
104        shape: ShapeVec,
105        placement: Placement,
106        layout_class: LayoutClass,
107        strides: StrideVec,
108        alignment_log2: Option<u8>,
109        physical_identity: InputPhysicalIdentity,
110    ) -> Self {
111        Self {
112            dtype,
113            shape,
114            placement,
115            layout_class,
116            strides,
117            alignment_log2,
118            backend_family: physical_identity.backend_family,
119            allocation_domain: physical_identity.allocation_domain,
120        }
121    }
122
123    /// Return the dtype component.
124    ///
125    /// # Examples
126    ///
127    /// ```
128    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
129    /// use tenferro_runtime::{DType, InputSignatureEntry, LayoutClass};
130    /// use tenferro_tensor::Placement;
131    ///
132    /// let entry = InputSignatureEntry::new(
133    ///     DType::Bool,
134    ///     [1_usize].into_iter().collect(),
135    ///     Placement::default(),
136    ///     LayoutClass::new("tenferro.layout.strided")?,
137    ///     [1_isize].into_iter().collect(),
138    ///     None,
139    /// )?;
140    /// assert_eq!(entry.dtype(), DType::Bool);
141    /// # Ok(())
142    /// # }
143    /// ```
144    pub fn dtype(&self) -> DType {
145        self.dtype
146    }
147
148    /// Return the shape component.
149    ///
150    /// # Examples
151    ///
152    /// ```
153    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
154    /// use tenferro_runtime::{DType, InputSignatureEntry, LayoutClass};
155    /// use tenferro_tensor::Placement;
156    ///
157    /// let entry = InputSignatureEntry::new(
158    ///     DType::F64,
159    ///     [2_usize, 3].into_iter().collect(),
160    ///     Placement::default(),
161    ///     LayoutClass::new("tenferro.layout.strided")?,
162    ///     [1_isize, 2].into_iter().collect(),
163    ///     None,
164    /// )?;
165    /// assert_eq!(entry.shape(), &[2, 3]);
166    /// # Ok(())
167    /// # }
168    /// ```
169    pub fn shape(&self) -> &[usize] {
170        &self.shape
171    }
172
173    /// Return the placement metadata component.
174    ///
175    /// # Examples
176    ///
177    /// ```
178    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
179    /// use tenferro_runtime::{DType, InputSignatureEntry, LayoutClass};
180    /// use tenferro_tensor::{MemoryKind, Placement};
181    ///
182    /// let entry = InputSignatureEntry::new(
183    ///     DType::F64,
184    ///     [1_usize].into_iter().collect(),
185    ///     Placement::default(),
186    ///     LayoutClass::new("tenferro.layout.strided")?,
187    ///     [1_isize].into_iter().collect(),
188    ///     None,
189    /// )?;
190    /// assert_eq!(entry.placement().memory_kind, MemoryKind::UnpinnedHost);
191    /// # Ok(())
192    /// # }
193    /// ```
194    pub fn placement(&self) -> &Placement {
195        &self.placement
196    }
197
198    /// Return the layout class component.
199    ///
200    /// # Examples
201    ///
202    /// ```
203    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
204    /// use tenferro_runtime::{DType, InputSignatureEntry, LayoutClass};
205    /// use tenferro_tensor::Placement;
206    ///
207    /// let layout = LayoutClass::new("tenferro.layout.strided")?;
208    /// let entry = InputSignatureEntry::new(
209    ///     DType::F64,
210    ///     [1_usize].into_iter().collect(),
211    ///     Placement::default(),
212    ///     layout.clone(),
213    ///     [1_isize].into_iter().collect(),
214    ///     None,
215    /// )?;
216    /// assert_eq!(entry.layout_class(), &layout);
217    /// # Ok(())
218    /// # }
219    /// ```
220    pub fn layout_class(&self) -> &LayoutClass {
221        &self.layout_class
222    }
223
224    /// Return the stride metadata component.
225    ///
226    /// # Examples
227    ///
228    /// ```
229    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
230    /// use tenferro_runtime::{DType, InputSignatureEntry, LayoutClass};
231    /// use tenferro_tensor::Placement;
232    ///
233    /// let entry = InputSignatureEntry::new(
234    ///     DType::F64,
235    ///     [2_usize].into_iter().collect(),
236    ///     Placement::default(),
237    ///     LayoutClass::new("tenferro.layout.strided")?,
238    ///     [2_isize].into_iter().collect(),
239    ///     None,
240    /// )?;
241    /// assert_eq!(entry.strides(), &[2]);
242    /// # Ok(())
243    /// # }
244    /// ```
245    pub fn strides(&self) -> &[isize] {
246        &self.strides
247    }
248
249    /// Return the known alignment class, if available.
250    ///
251    /// # Examples
252    ///
253    /// ```
254    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
255    /// use tenferro_runtime::{DType, InputSignatureEntry, LayoutClass};
256    /// use tenferro_tensor::Placement;
257    ///
258    /// let entry = InputSignatureEntry::new(
259    ///     DType::F64,
260    ///     [1_usize].into_iter().collect(),
261    ///     Placement::default(),
262    ///     LayoutClass::new("tenferro.layout.strided")?,
263    ///     [1_isize].into_iter().collect(),
264    ///     Some(3),
265    /// )?;
266    /// assert_eq!(entry.alignment_log2(), Some(3));
267    /// # Ok(())
268    /// # }
269    /// ```
270    pub fn alignment_log2(&self) -> Option<u8> {
271        self.alignment_log2
272    }
273
274    pub(super) fn backend_family(&self) -> Option<&'static str> {
275        self.backend_family
276    }
277
278    pub(super) fn allocation_domain(&self) -> Option<AllocationDomainId> {
279        self.allocation_domain
280    }
281
282    pub(crate) fn logical_retained_bytes(&self) -> Option<usize> {
283        checked_sum([
284            spilled_bytes::<usize>(self.shape.spilled(), self.shape.len())?,
285            spilled_bytes::<isize>(self.strides.spilled(), self.strides.len())?,
286        ])
287    }
288}
289
290/// Value-free signature of all tensor inputs for one prepare request.
291///
292/// # Examples
293///
294/// ```
295/// use tenferro_runtime::InputSignature;
296///
297/// let signature = InputSignature::new(Vec::new());
298/// assert!(signature.entries().is_empty());
299/// ```
300#[derive(Clone, Debug, Eq, Hash, PartialEq)]
301pub struct InputSignature {
302    entries: Vec<InputSignatureEntry>,
303}
304
305impl InputSignature {
306    /// Build a signature from already prepared entries.
307    ///
308    /// # Examples
309    ///
310    /// ```
311    /// use tenferro_runtime::InputSignature;
312    ///
313    /// let signature = InputSignature::new(Vec::new());
314    /// assert_eq!(signature.entries().len(), 0);
315    /// ```
316    pub fn new(entries: Vec<InputSignatureEntry>) -> Self {
317        Self { entries }
318    }
319
320    /// Build a value-free signature from borrowed tensor reads.
321    ///
322    /// # Examples
323    ///
324    /// ```
325    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
326    /// use tenferro_runtime::{InputSignature, TensorRead, Tensor};
327    ///
328    /// let tensor = Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0])?;
329    /// let signature = InputSignature::from_reads(&[TensorRead::from_tensor(&tensor)])?;
330    /// assert_eq!(signature.entries()[0].shape(), &[2]);
331    /// # Ok(())
332    /// # }
333    /// ```
334    ///
335    /// # Errors
336    ///
337    /// Returns [`PrepareError::InputSignature`] with the original typed tensor
338    /// metadata error when shape, stride, or compactness metadata cannot be read.
339    pub fn from_reads(reads: &[TensorRead<'_>]) -> Result<Self, PrepareError> {
340        let mut entries = Vec::with_capacity(reads.len());
341        for (input, read) in reads.iter().enumerate() {
342            let strides = read
343                .strides()
344                .map_err(|source| PrepareError::InputSignature {
345                    source: InputSignatureError::TensorMetadata { input, source },
346                })?;
347            let compact =
348                read.is_col_major_contiguous()
349                    .map_err(|source| PrepareError::InputSignature {
350                        source: InputSignatureError::TensorMetadata { input, source },
351                    })?;
352            let shape = read.shape().iter().copied().collect();
353            entries.push(InputSignatureEntry::from_validated_metadata(
354                read.dtype(),
355                shape,
356                read_placement(read),
357                layout_class(compact),
358                strides.into_iter().collect(),
359                read_alignment_log2(read),
360                InputPhysicalIdentity {
361                    backend_family: read.backend_family(),
362                    allocation_domain: read.allocation_domain(),
363                },
364            ));
365        }
366        Ok(Self { entries })
367    }
368
369    /// Return the per-input entries.
370    ///
371    /// # Examples
372    ///
373    /// ```
374    /// use tenferro_runtime::InputSignature;
375    ///
376    /// assert!(InputSignature::new(Vec::new()).entries().is_empty());
377    /// ```
378    pub fn entries(&self) -> &[InputSignatureEntry] {
379        &self.entries
380    }
381
382    pub(crate) fn logical_retained_bytes(&self) -> Option<usize> {
383        checked_sum([
384            self.entries
385                .len()
386                .checked_mul(size_of::<InputSignatureEntry>())?,
387            checked_sum_options(
388                self.entries
389                    .iter()
390                    .map(InputSignatureEntry::logical_retained_bytes),
391            )?,
392        ])
393    }
394}
395
396fn spilled_bytes<T>(spilled: bool, len: usize) -> Option<usize> {
397    if spilled {
398        len.checked_mul(size_of::<T>())
399    } else {
400        Some(0)
401    }
402}
403
404fn checked_sum(values: impl IntoIterator<Item = usize>) -> Option<usize> {
405    values
406        .into_iter()
407        .try_fold(0usize, |sum, value| sum.checked_add(value))
408}
409
410fn checked_sum_options(values: impl IntoIterator<Item = Option<usize>>) -> Option<usize> {
411    values
412        .into_iter()
413        .try_fold(0usize, |sum, value| sum.checked_add(value?))
414}
415
416fn validate_entry(
417    shape: &[usize],
418    strides: &[isize],
419    alignment_log2: Option<u8>,
420) -> Result<(), InputSignatureError> {
421    if shape.len() != strides.len() {
422        return Err(InputSignatureError::ShapeStrideRankMismatch {
423            rank: shape.len(),
424            stride_count: strides.len(),
425        });
426    }
427    if let Some(alignment_log2) = alignment_log2 {
428        if u32::from(alignment_log2) >= usize::BITS {
429            return Err(InputSignatureError::InvalidAlignmentClass { alignment_log2 });
430        }
431    }
432    Ok(())
433}
434
435pub(super) fn read_placement(read: &TensorRead<'_>) -> Placement {
436    match read {
437        TensorRead::Tensor(tensor) => tensor.placement().clone(),
438        TensorRead::View(view) => view_placement(view),
439    }
440}
441
442fn view_placement(view: &TensorView<'_>) -> Placement {
443    match view {
444        TensorView::F32(view) => view.placement().clone(),
445        TensorView::F64(view) => view.placement().clone(),
446        TensorView::I32(view) => view.placement().clone(),
447        TensorView::I64(view) => view.placement().clone(),
448        TensorView::Bool(view) => view.placement().clone(),
449        TensorView::C32(view) => view.placement().clone(),
450        TensorView::C64(view) => view.placement().clone(),
451    }
452}
453
454fn layout_class(compact: bool) -> LayoutClass {
455    let value = if compact {
456        COMPACT_COL_MAJOR_LAYOUT
457    } else {
458        STRIDED_LAYOUT
459    };
460    LayoutClass::runtime_created(value)
461}
462
463fn read_alignment_log2(read: &TensorRead<'_>) -> Option<u8> {
464    match read {
465        TensorRead::Tensor(tensor) => tensor_alignment_log2(tensor),
466        TensorRead::View(view) => view_alignment_log2(view),
467    }
468}
469
470fn tensor_alignment_log2(tensor: &Tensor) -> Option<u8> {
471    match tensor {
472        Tensor::F32(tensor) => typed_tensor_alignment_log2(tensor),
473        Tensor::F64(tensor) => typed_tensor_alignment_log2(tensor),
474        Tensor::I32(tensor) => typed_tensor_alignment_log2(tensor),
475        Tensor::I64(tensor) => typed_tensor_alignment_log2(tensor),
476        Tensor::Bool(tensor) => typed_tensor_alignment_log2(tensor),
477        Tensor::C32(tensor) => typed_tensor_alignment_log2(tensor),
478        Tensor::C64(tensor) => typed_tensor_alignment_log2(tensor),
479    }
480}
481
482fn typed_tensor_alignment_log2<T: TensorScalar>(tensor: &TypedTensor<T>) -> Option<u8> {
483    if tensor.buffer().is_backend() {
484        return None;
485    }
486    if shape_is_empty(tensor.shape()) {
487        return Some(type_alignment_log2::<T>());
488    }
489    tensor
490        .host_data()
491        .ok()
492        .map(|data| pointer_alignment_log2::<T>(data.as_ptr()))
493}
494
495fn view_alignment_log2(view: &TensorView<'_>) -> Option<u8> {
496    match view {
497        TensorView::F32(view) => typed_view_alignment_log2(view),
498        TensorView::F64(view) => typed_view_alignment_log2(view),
499        TensorView::I32(view) => typed_view_alignment_log2(view),
500        TensorView::I64(view) => typed_view_alignment_log2(view),
501        TensorView::Bool(view) => typed_view_alignment_log2(view),
502        TensorView::C32(view) => typed_view_alignment_log2(view),
503        TensorView::C64(view) => typed_view_alignment_log2(view),
504    }
505}
506
507fn typed_view_alignment_log2<T: 'static>(view: &TypedTensorView<'_, T>) -> Option<u8> {
508    if view.backend_buffer().is_some() {
509        return None;
510    }
511    if shape_is_empty(view.shape()) {
512        return Some(type_alignment_log2::<T>());
513    }
514    view.host_storage().ok().map(|data| {
515        let pointer = data.as_ptr().wrapping_offset(view.offset());
516        pointer_alignment_log2::<T>(pointer)
517    })
518}
519
520fn shape_is_empty(shape: &[usize]) -> bool {
521    shape.contains(&0)
522}
523
524fn type_alignment_log2<T>() -> u8 {
525    align_of::<T>().trailing_zeros().min(usize::BITS - 1) as u8
526}
527
528fn pointer_alignment_log2<T>(pointer: *const T) -> u8 {
529    (pointer as usize)
530        .trailing_zeros()
531        .min(align_of::<T>().trailing_zeros())
532        .min(usize::BITS - 1) as u8
533}
tenferro_runtime/runtime/signature.rs

tenferro_runtime/runtime/
signature.rs
