somatize_core/
build.rs

1//! Declaring a graph as an expression, instead of by calls.
2//!
3//! ```ignore
4//! let (graph, catalog, placement, memory) = (node("source", Add(1.0))
5//!     >> (node("left", Add(10.0)).on(Device::Cuda(0)) | node("right", Add(100.0)))
6//!     >> node("join", Mean))
7//! .somatize()?;
8//! ```
9//!
10//! `>>` chains and `|` opens branches, the same syntax as the Python DSL. In
11//! Rust it falls out of implementing [`std::ops::Shr`] and
12//! [`std::ops::BitOr`] on a type of our own; no macro needed.
13//!
14//! A [`Wire`] is a half-declared graph: where you enter (`heads`) and where you
15//! leave (`terminals`), which is all it takes to glue another one on. Nothing is
16//! materialized until [`Wire::somatize`], so joining two pieces concatenates
17//! lists rather than merging graphs.
18
19use crate::{Catalog, Device, Graph, GraphError, Host, Memory, Node, NodeId, Placement};
20use std::ops::{BitOr, Shr};
21use std::sync::Arc;
22
23/// A half-declared graph.
24pub struct Wire {
25    parts: Result<Parts, GraphError>,
26}
27
28struct Parts {
29    nodes: Vec<(NodeId, Arc<dyn Node>)>,
30    edges: Vec<(NodeId, NodeId)>,
31    heads: Vec<NodeId>,
32    terminals: Vec<NodeId>,
33    /// The ones that already have a device. An id appears at most once.
34    devices: Vec<(NodeId, Device)>,
35    /// The ones that already have a host. Separate from the devices so that
36    /// `.on(...)` does not shadow an inner `.at(...)`, or the other way round.
37    hosts: Vec<(NodeId, Host)>,
38    /// What implements each one. Filled where the concrete type is still known,
39    /// which is [`node`] and nowhere else: from there on it is an `Arc<dyn
40    /// Node>` and the name is gone.
41    identities: Vec<(NodeId, String)>,
42    /// The ones settled, each with the digest of the state they are settled at —
43    /// never one here, because hashing weights is torch's job and this is the
44    /// core.
45    frozen: Vec<(NodeId, Option<String>)>,
46    /// The ones worth keeping, each with its salt — likewise never one here:
47    /// telling apart two runs the key cannot is a knob for whoever runs them,
48    /// and it is [`Memory::cache`] for anyone who wants it.
49    cached: Vec<(NodeId, Option<String>)>,
50    /// The ones that map over the items of their input.
51    mapped: Vec<NodeId>,
52}
53
54/// A lone node, named after its type: this is the last place that knows it, and
55/// what a node is called is half of the key its output is kept under.
56pub fn node<N: Node + 'static>(id: impl Into<NodeId>, implementation: N) -> Wire {
57    single(
58        id.into(),
59        std::any::type_name::<N>(),
60        Arc::new(implementation),
61    )
62}
63
64fn single(id: NodeId, identity: &str, implementation: Arc<dyn Node>) -> Wire {
65    Wire {
66        parts: Ok(Parts {
67            nodes: vec![(id.clone(), implementation)],
68            edges: Vec::new(),
69            heads: vec![id.clone()],
70            terminals: vec![id.clone()],
71            devices: Vec::new(),
72            hosts: Vec::new(),
73            identities: vec![(id, identity.to_string())],
74            frozen: Vec::new(),
75            cached: Vec::new(),
76            mapped: Vec::new(),
77        }),
78    }
79}
80
81impl Shr for Wire {
82    type Output = Wire;
83
84    /// `a >> b`: everything leaving `a` enters everything starting `b`.
85    fn shr(self, next: Wire) -> Wire {
86        combine(self, next, |left, right| Parts {
87            edges: left
88                .terminals
89                .iter()
90                .flat_map(|from| right.heads.iter().map(|to| (from.clone(), to.clone())))
91                .chain(left.edges)
92                .chain(right.edges)
93                .collect(),
94            nodes: left.nodes.into_iter().chain(right.nodes).collect(),
95            heads: left.heads,
96            terminals: right.terminals,
97            devices: left.devices.into_iter().chain(right.devices).collect(),
98            hosts: left.hosts.into_iter().chain(right.hosts).collect(),
99            identities: left
100                .identities
101                .into_iter()
102                .chain(right.identities)
103                .collect(),
104            frozen: left.frozen.into_iter().chain(right.frozen).collect(),
105            cached: left.cached.into_iter().chain(right.cached).collect(),
106            mapped: left.mapped.into_iter().chain(right.mapped).collect(),
107        })
108    }
109}
110
111impl BitOr for Wire {
112    type Output = Wire;
113
114    /// `a | b`: two branches that do not touch. Whatever comes in reaches both,
115    /// and whatever leaves either one leaves here.
116    fn bitor(self, other: Wire) -> Wire {
117        combine(self, other, |left, right| Parts {
118            nodes: left.nodes.into_iter().chain(right.nodes).collect(),
119            edges: left.edges.into_iter().chain(right.edges).collect(),
120            heads: left.heads.into_iter().chain(right.heads).collect(),
121            terminals: left.terminals.into_iter().chain(right.terminals).collect(),
122            devices: left.devices.into_iter().chain(right.devices).collect(),
123            hosts: left.hosts.into_iter().chain(right.hosts).collect(),
124            identities: left
125                .identities
126                .into_iter()
127                .chain(right.identities)
128                .collect(),
129            frozen: left.frozen.into_iter().chain(right.frozen).collect(),
130            cached: left.cached.into_iter().chain(right.cached).collect(),
131            mapped: left.mapped.into_iter().chain(right.mapped).collect(),
132        })
133    }
134}
135
136/// Joins two pieces, keeping the first failure if either carries one.
137fn combine(left: Wire, right: Wire, join: impl FnOnce(Parts, Parts) -> Parts) -> Wire {
138    Wire {
139        parts: match (left.parts, right.parts) {
140            (Ok(left), Ok(right)) => Ok(join(left, right)),
141            (Err(e), _) | (_, Err(e)) => Err(e),
142        },
143    }
144}
145
146/// Gives `what` to the nodes that did not already have something of that half:
147/// the "innermost one wins" rule, written once for both.
148fn fill<T: Clone>(nodes: &[(NodeId, Arc<dyn Node>)], placed: &mut Vec<(NodeId, T)>, what: T) {
149    let unplaced: Vec<NodeId> = nodes
150        .iter()
151        .map(|(id, _)| id)
152        .filter(|id| !placed.iter().any(|(already, _)| already == *id))
153        .cloned()
154        .collect();
155    placed.extend(unplaced.into_iter().map(|id| (id, what.clone())));
156}
157
158impl Wire {
159    /// This whole piece on one device. The innermost one wins:
160    /// `(a.on(Cuda(0)) >> b).on(Cuda(1))` leaves `a` on 0 and `b` on 1.
161    pub fn on(self, device: Device) -> Wire {
162        Wire {
163            parts: self.parts.map(|mut parts| {
164                fill(&parts.nodes, &mut parts.devices, device);
165                parts
166            }),
167        }
168    }
169
170    /// This whole piece on one host, likewise, and **independent** of the
171    /// device: the two can be written in any order.
172    pub fn at(self, host: impl Into<Host>) -> Wire {
173        let host = host.into();
174        Wire {
175            parts: self.parts.map(|mut parts| {
176                fill(&parts.nodes, &mut parts.hosts, host);
177                parts
178            }),
179        }
180    }
181
182    /// This whole piece settled: its state does not change while the graph
183    /// runs, innermost first. Only the half the core can hold — whoever knows
184    /// what a gradient is says it again with the digest. See [`Memory::freeze`].
185    pub fn frozen(self) -> Wire {
186        Wire {
187            parts: self.parts.map(|mut parts| {
188                fill(&parts.nodes, &mut parts.frozen, None);
189                parts
190            }),
191        }
192    }
193
194    /// This whole piece worth keeping: what each node produces is looked up
195    /// before being computed and kept after. Whether it can honestly be kept is
196    /// [`cacheable`](crate::cacheable)'s question, asked before running.
197    pub fn cached(self) -> Wire {
198        Wire {
199            parts: self.parts.map(|mut parts| {
200                fill(&parts.nodes, &mut parts.cached, None);
201                parts
202            }),
203        }
204    }
205
206    /// This whole piece maps over the items of its input: a list in, a list as
207    /// long out. What gives a cache the grain of an **item**.
208    pub fn mapped(self) -> Wire {
209        Wire {
210            parts: self.parts.map(|mut parts| {
211                parts
212                    .mapped
213                    .extend(parts.nodes.iter().map(|(id, _)| id.clone()));
214                parts
215            }),
216        }
217    }
218
219    /// Materializes what was declared: the structure, the store, the placement
220    /// and what is remembered, none containing the others. Fails on a repeated
221    /// id, above all.
222    pub fn somatize(self) -> Result<(Graph, Catalog, Placement, Memory), GraphError> {
223        let parts = self.parts?;
224        let mut graph = Graph::new();
225        let mut catalog = Catalog::new();
226        let mut placement = Placement::new();
227        let mut memory = Memory::new();
228
229        for (id, implementation) in parts.nodes {
230            graph.add_node(id.clone())?;
231            catalog.insert(id, implementation);
232        }
233        for (from, to) in parts.edges {
234            graph.add_edge(from, to)?;
235        }
236        for (id, device) in parts.devices {
237            placement.place(id, device);
238        }
239        for (id, host) in parts.hosts {
240            placement.place_at(id, host);
241        }
242        for (id, what) in parts.identities {
243            memory.identify(id, what);
244        }
245        for (id, state) in parts.frozen {
246            memory.freeze(id, state);
247        }
248        for (id, salt) in parts.cached {
249            memory.cache(id, salt);
250        }
251        for id in parts.mapped {
252            memory.map(id);
253        }
254        Ok((graph, catalog, placement, memory))
255    }
256}
somatize_core/build.rs

somatize_core/
build.rs
