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>,
}
#[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)
})
}
pub fn len(&self) -> usize {
self.entries.len()
}
pub fn is_empty(&self) -> bool {
self.entries.is_empty()
}
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())
}
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()
}
}