tenferro-runtime 0.8.0

Concrete tensor helpers, traced tensors, graph execution, and extension runtime support for tenferro.
Documentation
use std::fmt;
use std::sync::Arc;

use crate::checkpoint::RetainedValue;

use super::value::ProgramBuilderNonce;
use super::{BindingKey, ProgramValue};

pub(crate) struct PendingBinding {
    pub(crate) key: BindingKey,
    pub(crate) input: ProgramValue,
    pub(crate) tensor: Arc<RetainedValue>,
}

struct ProgramBinding {
    key: BindingKey,
    tensor: Arc<RetainedValue>,
}

/// Immutable tensor defaults and large constants kept outside semantic structure.
#[derive(Clone)]
pub struct ProgramBindings {
    owner: ProgramBuilderNonce,
    entries: Arc<[ProgramBinding]>,
}

impl ProgramBindings {
    pub(crate) fn freeze(owner: ProgramBuilderNonce, pending: Vec<PendingBinding>) -> Self {
        let mut entries: Vec<_> = pending
            .into_iter()
            .map(|binding| ProgramBinding {
                key: binding.key,
                tensor: binding.tensor,
            })
            .collect();
        entries.sort_unstable_by_key(|entry| entry.key.slot);
        Self {
            owner,
            entries: entries.into(),
        }
    }

    pub(crate) fn belongs_to(&self, owner: ProgramBuilderNonce) -> bool {
        self.owner == owner
    }

    pub(crate) fn tensor_for_input(&self, input: ProgramValue) -> Option<Arc<RetainedValue>> {
        self.tensor_ref_for_input(input).map(Arc::clone)
    }

    pub(crate) fn tensor_ref_for_input(&self, input: ProgramValue) -> Option<&Arc<RetainedValue>> {
        if input.owner != self.owner {
            return None;
        }
        self.entries
            .binary_search_by_key(&input.slot, |entry| entry.key.slot)
            .ok()
            .and_then(|index| self.entries.get(index))
            .map(|entry| &entry.tensor)
    }

    pub(crate) fn bound_inputs(&self) -> impl ExactSizeIterator<Item = ProgramValue> + '_ {
        self.entries
            .iter()
            .map(|entry| ProgramValue::new(entry.key.slot, self.owner))
    }

    #[allow(dead_code, reason = "wired into GraphCompiler in Phase 3 A3")]
    pub(crate) fn preserves_all(&self, source: &Self) -> bool {
        if self.entries.len() < source.entries.len() {
            return false;
        }
        let mut matched = vec![false; self.entries.len()];
        source.entries.iter().all(|source_entry| {
            self.entries
                .iter()
                .enumerate()
                .find(|(index, entry)| {
                    !matched[*index] && Arc::ptr_eq(&entry.tensor, &source_entry.tensor)
                })
                .map(|(index, _)| {
                    matched[index] = true;
                })
                .is_some()
        })
    }

    #[allow(dead_code, reason = "wired into GraphCompiler in Phase 3 A3")]
    pub(crate) fn cache_exact_eq(&self, other: &Self) -> bool {
        self.entries.len() == other.entries.len()
            && self
                .entries
                .iter()
                .zip(other.entries.iter())
                .all(|(left, right)| {
                    left.key.slot == right.key.slot && Arc::ptr_eq(&left.tensor, &right.tensor)
                })
    }

    /// Return the number of tensor bindings.
    pub fn len(&self) -> usize {
        self.entries.len()
    }

    /// Return whether no tensor bindings are present.
    pub fn is_empty(&self) -> bool {
        self.entries.is_empty()
    }

    /// Borrow a tensor by opaque binding key.
    pub fn get(&self, key: BindingKey) -> Option<&RetainedValue> {
        if key.owner != self.owner {
            return None;
        }
        self.entries
            .binary_search_by_key(&key.slot, |entry| entry.key.slot)
            .ok()
            .and_then(|index| self.entries.get(index))
            .filter(|entry| entry.key == key)
            .map(|entry| entry.tensor.as_ref())
    }

    /// Iterate over ordered binding keys and borrowed move-only values.
    pub fn iter(&self) -> impl ExactSizeIterator<Item = (BindingKey, &RetainedValue)> + '_ {
        self.entries
            .iter()
            .map(|entry| (entry.key, entry.tensor.as_ref()))
    }
}

impl fmt::Debug for ProgramBindings {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("ProgramBindings")
            .field("len", &self.entries.len())
            .finish()
    }
}