tenferro_runtime/program/
bindings.rs1use std::fmt;
2use std::sync::Arc;
3
4use crate::checkpoint::RetainedValue;
5
6use super::value::ProgramBuilderNonce;
7use super::{BindingKey, ProgramValue};
8
9#[derive(Clone)]
10pub(crate) struct PendingBinding {
11 pub(crate) key: BindingKey,
12 pub(crate) input: ProgramValue,
13 pub(crate) tensor: Arc<RetainedValue>,
14}
15
16struct ProgramBinding {
17 key: BindingKey,
18 tensor: Arc<RetainedValue>,
19}
20
21#[derive(Clone)]
23pub struct ProgramBindings {
24 owner: ProgramBuilderNonce,
25 entries: Arc<[ProgramBinding]>,
26}
27
28impl ProgramBindings {
29 pub(crate) fn freeze(owner: ProgramBuilderNonce, pending: Vec<PendingBinding>) -> Self {
30 let mut entries: Vec<_> = pending
31 .into_iter()
32 .map(|binding| ProgramBinding {
33 key: binding.key,
34 tensor: binding.tensor,
35 })
36 .collect();
37 entries.sort_unstable_by_key(|entry| entry.key.slot);
38 Self {
39 owner,
40 entries: entries.into(),
41 }
42 }
43
44 pub(crate) fn belongs_to(&self, owner: ProgramBuilderNonce) -> bool {
45 self.owner == owner
46 }
47
48 pub(crate) fn tensor_for_input(&self, input: ProgramValue) -> Option<Arc<RetainedValue>> {
49 self.tensor_ref_for_input(input).map(Arc::clone)
50 }
51
52 pub(crate) fn tensor_ref_for_input(&self, input: ProgramValue) -> Option<&Arc<RetainedValue>> {
53 if input.owner != self.owner {
54 return None;
55 }
56 self.entries
57 .binary_search_by_key(&input.slot, |entry| entry.key.slot)
58 .ok()
59 .and_then(|index| self.entries.get(index))
60 .map(|entry| &entry.tensor)
61 }
62
63 pub(crate) fn bound_inputs(&self) -> impl ExactSizeIterator<Item = ProgramValue> + '_ {
64 self.entries
65 .iter()
66 .map(|entry| ProgramValue::new(entry.key.slot, self.owner))
67 }
68
69 #[allow(dead_code, reason = "wired into GraphCompiler in Phase 3 A3")]
70 pub(crate) fn preserves_all(&self, source: &Self) -> bool {
71 if self.entries.len() < source.entries.len() {
72 return false;
73 }
74 let mut matched = vec![false; self.entries.len()];
75 source.entries.iter().all(|source_entry| {
76 self.entries
77 .iter()
78 .enumerate()
79 .find(|(index, entry)| {
80 !matched[*index] && Arc::ptr_eq(&entry.tensor, &source_entry.tensor)
81 })
82 .map(|(index, _)| {
83 matched[index] = true;
84 })
85 .is_some()
86 })
87 }
88
89 #[allow(dead_code, reason = "wired into GraphCompiler in Phase 3 A3")]
90 pub(crate) fn cache_exact_eq(&self, other: &Self) -> bool {
91 self.entries.len() == other.entries.len()
92 && self
93 .entries
94 .iter()
95 .zip(other.entries.iter())
96 .all(|(left, right)| {
97 left.key.slot == right.key.slot && Arc::ptr_eq(&left.tensor, &right.tensor)
98 })
99 }
100
101 pub fn len(&self) -> usize {
103 self.entries.len()
104 }
105
106 pub fn is_empty(&self) -> bool {
108 self.entries.is_empty()
109 }
110
111 pub fn get(&self, key: BindingKey) -> Option<&RetainedValue> {
113 if key.owner != self.owner {
114 return None;
115 }
116 self.entries
117 .binary_search_by_key(&key.slot, |entry| entry.key.slot)
118 .ok()
119 .and_then(|index| self.entries.get(index))
120 .filter(|entry| entry.key == key)
121 .map(|entry| entry.tensor.as_ref())
122 }
123
124 pub fn iter(&self) -> impl ExactSizeIterator<Item = (BindingKey, &RetainedValue)> + '_ {
126 self.entries
127 .iter()
128 .map(|entry| (entry.key, entry.tensor.as_ref()))
129 }
130}
131
132impl fmt::Debug for ProgramBindings {
133 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
134 formatter
135 .debug_struct("ProgramBindings")
136 .field("len", &self.entries.len())
137 .finish()
138 }
139}