use std::borrow::Borrow;
use std::fmt::{Debug, Display};
use multithread::Executor;
use crate::internal::*;
use crate::model::{Fact, Graph, OutletId};
use crate::ops::konst::Const;
use crate::runtime::RunOptions;
use self::order::{build_flush_list, eval_order_for_nodes, eval_order_opt_ram_for_nodes};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct SessionId(u64);
static NEXT_SESSION_ID: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
impl SessionId {
pub const NONE: SessionId = SessionId(u64::MAX);
fn next() -> SessionId {
SessionId(NEXT_SESSION_ID.fetch_add(1, std::sync::atomic::Ordering::Relaxed))
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct LaneId(pub usize);
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Seating {
lanes: Vec<LaneId>,
max_lanes: usize,
}
impl Seating {
pub fn new(max_lanes: usize, lanes: impl IntoIterator<Item = LaneId>) -> TractResult<Seating> {
let lanes: Vec<LaneId> = lanes.into_iter().collect();
for (seat, lane) in lanes.iter().enumerate() {
ensure!(lane.0 < max_lanes, "Seat {seat} takes lane {} of {max_lanes}", lane.0);
ensure!(!lanes[..seat].contains(lane), "Lane {} takes two seats in one turn", lane.0);
}
Ok(Seating { lanes, max_lanes })
}
pub fn single() -> Seating {
Seating { lanes: vec![LaneId(0)], max_lanes: 1 }
}
pub fn max_lanes(&self) -> usize {
self.max_lanes
}
pub fn lanes(&self) -> &[LaneId] {
&self.lanes
}
pub fn occupancy(&self) -> usize {
self.lanes.len()
}
pub fn address(&self, ix: usize) -> (Option<usize>, Option<usize>) {
if self.max_lanes == 1 { (None, None) } else { (Some(ix), Some(self.lanes[ix].0)) }
}
}
pub type TurnShared = anymap3::Map<dyn std::any::Any + Send>;
#[derive(Debug, Clone, Copy)]
pub struct EvalContext<'a> {
pub session: SessionId,
pub node_id: usize,
pub symbols: &'a SymbolValues,
pub scenario: Option<usize>,
pub shared: Option<&'a TurnShared>,
pub seating: &'a Seating,
}
pub struct TurnState {
pub resolved_symbols: SymbolValues,
pub scenario: Option<usize>,
pub values: Vec<Option<TVec<TValue>>>,
pub shared: TurnShared,
pub seating: Seating,
}
impl EvalContext<'static> {
pub fn out_of_plan() -> EvalContext<'static> {
static SYMBOLS: std::sync::OnceLock<SymbolValues> = std::sync::OnceLock::new();
static SEATING: std::sync::OnceLock<Seating> = std::sync::OnceLock::new();
EvalContext {
session: SessionId::NONE,
node_id: usize::MAX,
symbols: SYMBOLS.get_or_init(SymbolValues::default),
scenario: None,
shared: None,
seating: SEATING.get_or_init(Seating::single),
}
}
}
impl TurnState {
pub fn context(&self, session: SessionId, node_id: usize) -> EvalContext<'_> {
EvalContext {
session,
node_id,
symbols: &self.resolved_symbols,
scenario: self.scenario,
shared: Some(&self.shared),
seating: &self.seating,
}
}
}
impl Default for TurnState {
fn default() -> Self {
TurnState {
resolved_symbols: SymbolValues::default(),
scenario: None,
values: vec![],
shared: TurnShared::new(),
seating: Seating::single(),
}
}
}
impl Clone for TurnState {
fn clone(&self) -> Self {
TurnState {
resolved_symbols: self.resolved_symbols.clone(),
scenario: self.scenario,
values: vec![],
shared: TurnShared::new(),
seating: self.seating.clone(),
}
}
}
pub trait TurnStateHandler: Send + Sync + Debug {
fn before_plan_eval(&self, turn: &mut TurnState) -> TractResult<()>;
fn after_plan_eval(&self, turn: &mut TurnState) -> TractResult<()>;
}
impl Debug for TurnState {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "TurnState({:?})", self.resolved_symbols)
}
}
#[derive(Debug, Clone)]
pub struct SimplePlan<F, O>
where
F: Fact + Clone + 'static,
O: Debug + Display + AsRef<dyn Op> + AsMut<dyn Op> + Clone + 'static,
{
pub(crate) model: Arc<Graph<F, O>>,
outputs: Vec<OutletId>,
order: Vec<usize>,
flush_lists: Vec<TVec<usize>>,
has_unresolved_symbols: bool,
symbols: Vec<Symbol>,
executor: Option<Executor>,
turn_handler: Option<Arc<dyn TurnStateHandler + 'static>>,
const_values: Vec<(usize, TValue)>,
}
impl<F, O> SimplePlan<F, O>
where
F: Fact + Clone + 'static,
O: Debug + Display + AsRef<dyn Op> + AsMut<dyn Op> + Clone + 'static,
{
pub fn new(model: impl Into<Arc<Graph<F, O>>>) -> TractResult<Arc<SimplePlan<F, O>>> {
let model = model.into();
Self::build(model, &RunOptions::default()).map(Arc::new)
}
pub fn new_with_options(
model: impl Into<Arc<Graph<F, O>>>,
options: &RunOptions,
) -> TractResult<Arc<SimplePlan<F, O>>> {
let model = model.into();
Self::build(model, options).map(Arc::new)
}
#[deprecated]
pub fn new_for_output(
model: Graph<F, O>,
output: OutletId,
) -> TractResult<Arc<SimplePlan<F, O>>> {
#[allow(deprecated)]
Self::build_with_outputs_and_deps(model, &[output], &[], &RunOptions::default())
.map(Arc::new)
}
#[deprecated]
pub fn new_for_outputs(
model: impl Into<Arc<Graph<F, O>>>,
outputs: &[OutletId],
) -> TractResult<Arc<SimplePlan<F, O>>> {
#[allow(deprecated)]
Self::build_with_outputs_and_deps(model, outputs, &[], &RunOptions::default()).map(Arc::new)
}
pub fn with_turn_handler<H: TurnStateHandler + 'static>(mut self, turn_handler: H) -> Self {
self.turn_handler = Some(Arc::new(turn_handler));
self
}
#[deprecated]
pub fn new_for_outputs_and_deps(
model: impl Into<Arc<Graph<F, O>>>,
outputs: &[OutletId],
deps: &[(usize, usize)],
) -> TractResult<Arc<SimplePlan<F, O>>> {
#[allow(deprecated)]
Self::build_with_outputs_and_deps(model, outputs, deps, &RunOptions::default())
.map(Arc::new)
}
pub fn build(
model: impl Into<Arc<Graph<F, O>>>,
options: &RunOptions,
) -> TractResult<SimplePlan<F, O>> {
let model = model.into();
let outputs = model.outputs.clone();
#[allow(deprecated)]
Self::build_with_outputs_and_deps(model, &outputs, &[], options)
}
#[allow(clippy::mutable_key_type)]
fn symbol_deps(
nodes: &[Node<F, O>],
inputs: &[usize],
outputs: &[usize],
deps: &[(usize, usize)],
) -> TractResult<Vec<(usize, usize)>> {
fn tensor_value_symbols(t: &Tensor, acc: &mut std::collections::HashSet<Symbol>) {
if t.datum_type() == TDim::datum_type()
&& let Some(slice) =
t.try_as_plain_ram().ok().and_then(|v| v.as_slice::<TDim>().ok())
{
acc.extend(slice.iter().flat_map(|d| d.symbols()));
}
}
let base_order = eval_order_for_nodes(nodes, inputs, outputs, &[])?;
let mut definer: std::collections::HashMap<Symbol, usize> = Default::default();
for &id in inputs {
for output in &nodes[id].outputs {
if let Ok(fact) = output.fact.to_typed_fact() {
for sym in fact.shape.iter().flat_map(|d| d.symbols()) {
definer.entry(sym).or_insert(id);
}
}
}
}
for &id in &base_order {
let node = &nodes[id];
if !node.op.as_ref().mints_runtime_symbols() {
continue;
}
for output in &node.outputs {
if let Ok(fact) = output.fact.to_typed_fact() {
for sym in fact.shape.iter().flat_map(|d| d.symbols()) {
definer.entry(sym).or_insert(id);
}
}
}
}
let mut derived = vec![];
let mut seen: std::collections::HashSet<(usize, usize)> = deps.iter().copied().collect();
let mut definer_ancestors: std::collections::HashMap<
usize,
std::collections::HashSet<usize>,
> = Default::default();
for &id in &base_order {
let node = &nodes[id];
let mut mentioned: std::collections::HashSet<Symbol> = Default::default();
for output in &node.outputs {
if let Ok(fact) = output.fact.to_typed_fact() {
mentioned.extend(fact.shape.iter().flat_map(|d| d.symbols()));
for t in [fact.konst.as_deref(), fact.uniform.as_deref()].into_iter().flatten()
{
tensor_value_symbols(t, &mut mentioned);
}
mentioned.extend(fact.uniform_tdim.iter().flat_map(TDim::symbols));
}
}
let input_facts: Vec<Cow<'_, TypedFact>> = node
.inputs
.iter()
.filter_map(|i| nodes[i.node].outputs[i.slot].fact.to_typed_fact().ok())
.collect();
let input_facts: Vec<&TypedFact> = input_facts.iter().map(|f| f.as_ref()).collect();
mentioned.extend(node.op.as_ref().required_bound_symbols(&input_facts));
let mut mentioned: Vec<_> = mentioned.into_iter().collect();
mentioned.sort_unstable();
for sym in mentioned {
let Some(&d) = definer.get(&sym) else { continue };
if d == id || inputs.contains(&d) || !seen.insert((id, d)) {
continue;
}
if node.inputs.iter().any(|i| i.node == d) {
continue;
}
let ancestors = definer_ancestors.entry(d).or_insert_with(|| {
let mut stack: TVec<usize> = nodes[d].inputs.iter().map(|i| i.node).collect();
let mut set: std::collections::HashSet<usize> = Default::default();
while let Some(n) = stack.pop() {
if set.insert(n) {
stack.extend(nodes[n].inputs.iter().map(|i| i.node));
}
}
set
});
if !ancestors.contains(&id) {
derived.push((id, d));
}
}
}
Ok(derived)
}
#[deprecated]
pub fn build_with_outputs_and_deps(
model: impl Into<Arc<Graph<F, O>>>,
outputs: &[OutletId],
deps: &[(usize, usize)],
options: &RunOptions,
) -> TractResult<SimplePlan<F, O>> {
let model = model.into();
let inputs = model.input_outlets()?.iter().map(|n| n.node).collect::<Vec<usize>>();
let outputs_nodes = outputs.iter().map(|n| n.node).collect::<Vec<usize>>();
let mut all_deps = deps.to_vec();
all_deps.extend(Self::symbol_deps(model.nodes(), &inputs, &outputs_nodes, deps)?);
let mut order = if options.skip_order_opt_ram {
eval_order_for_nodes(model.nodes(), &inputs, &outputs_nodes, &all_deps)?
} else {
eval_order_opt_ram_for_nodes(model.nodes(), &inputs, &outputs_nodes, &all_deps)?
};
order.retain(|node| !model.node(*node).op_is::<Const>());
let flush_lists = build_flush_list(&*model, &order, outputs, |n| !n.op_is::<Const>());
#[allow(clippy::mutable_key_type)]
let mut symbols: std::collections::HashSet<Symbol> = Default::default();
for node in &model.nodes {
for output in &node.outputs {
if let Ok(fact) = output.fact.to_typed_fact() {
symbols.extend(fact.shape.iter().flat_map(|d| d.symbols()))
}
}
}
let const_values = model
.nodes
.iter()
.filter_map(|n| n.op_as::<Const>().map(|k| (n.id, k.val().clone().into_tvalue())))
.collect();
Ok(SimplePlan {
model,
order,
flush_lists,
outputs: outputs.to_vec(),
has_unresolved_symbols: !symbols.is_empty(),
symbols: symbols.into_iter().collect(),
executor: options.executor.clone(),
turn_handler: None,
const_values,
})
}
pub fn order_without_consts(&self) -> &[usize] {
&self.order
}
pub fn run(self: &Arc<Self>, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
let mut state = self.spawn()?;
state.run(inputs)
}
pub fn model(&self) -> &Graph<F, O> {
self.model.borrow()
}
pub fn spawn(self: &Arc<Self>) -> TractResult<SimpleState<F, O>> {
SimpleState::new(self)
}
}
#[derive(Debug)]
pub struct SimpleState<F, O>
where
F: Fact + Clone + 'static,
O: Debug + Display + AsRef<dyn Op> + AsMut<dyn Op> + Clone + 'static,
{
pub(crate) plan: Arc<SimplePlan<F, O>>,
pub op_states: Vec<Option<Box<dyn OpState>>>,
pub turn_state: TurnState,
session: SessionId,
}
impl<F, O> Clone for SimpleState<F, O>
where
F: Fact + Clone + 'static,
O: Debug + Display + AsRef<dyn Op> + AsMut<dyn Op> + Clone + 'static,
{
fn clone(&self) -> Self {
SimpleState {
plan: self.plan.clone(),
op_states: self.op_states.clone(),
turn_state: self.turn_state.clone(),
session: SessionId::next(),
}
}
}
impl<F, O> Drop for SimpleState<F, O>
where
F: Fact + Clone + 'static,
O: Debug + Display + AsRef<dyn Op> + AsMut<dyn Op> + Clone + 'static,
{
fn drop(&mut self) {
for (ix, node) in self.plan.model.nodes.iter().enumerate() {
node.op().drop_session(self.session, ix);
}
}
}
impl<F, O> SimpleState<F, O>
where
F: Fact + Clone + 'static,
O: Debug + Display + AsRef<dyn Op> + AsMut<dyn Op> + Clone + 'static,
{
pub fn new(plan: &Arc<SimplePlan<F, O>>) -> TractResult<SimpleState<F, O>> {
let plan = Arc::clone(plan);
let turn = TurnState::default();
let model = plan.model();
let states: Vec<Option<Box<dyn OpState>>> = vec![None; model.nodes.len()];
let mut state =
SimpleState { plan, op_states: states, turn_state: turn, session: SessionId::next() };
state.reset_op_states()?;
Ok(state)
}
pub fn new_from_inputs(
plan: &Arc<SimplePlan<F, O>>,
inputs: TVec<TValue>,
) -> TractResult<SimpleState<F, O>> {
let mut state = SimpleState::new(plan)?;
state.set_inputs(inputs)?;
state.resolve_symbols_with_states()?;
Ok(state)
}
fn ready_turn(&mut self) {
if self.turn_state.values.len() == 0 {
self.turn_state.values = vec![None; self.plan.model.nodes().len()];
for (id, value) in &self.plan.const_values {
self.turn_state.values[*id] = Some(tvec!(value.clone()));
}
}
}
pub fn reset_turn(&mut self) -> TractResult<()> {
self.reset_turn_keep_symbols();
self.turn_state.resolved_symbols = SymbolValues::default();
Ok(())
}
pub(crate) fn reset_turn_keep_symbols(&mut self) {
for node in &self.plan.order {
self.turn_state.values[*node] = None;
}
}
pub(crate) fn clear_resolved_symbols(&mut self) {
self.turn_state.resolved_symbols = SymbolValues::default();
self.turn_state.scenario = None;
}
pub fn seat(&mut self, seating: Seating) {
self.turn_state.seating = seating;
}
pub fn reset_lanes(&mut self, lanes: &[LaneId]) -> TractResult<()> {
for op_state in self.op_states.iter_mut().flatten() {
op_state.reset_lanes(lanes)?;
}
Ok(())
}
fn reset_op_states(&mut self) -> TractResult<()> {
let &mut SimpleState {
ref plan, ref turn_state, op_states: ref mut states, session, ..
} = self;
for (ix, n) in plan.model.nodes.iter().enumerate() {
states[ix] = n.op().state(&turn_state.context(session, ix))?;
}
Ok(())
}
pub(crate) fn resolve_symbols_with_states(&mut self) -> TractResult<()> {
for state in self
.op_states
.iter_mut()
.filter_map(Option::as_mut)
.filter(|s| s.has_init_tensor_fact())
{
state.resolve_symbols(&mut self.turn_state)?;
}
Ok(())
}
pub fn run(&mut self, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
self.run_plan_with_eval(inputs, self::eval)
}
pub fn exec(&mut self) -> TractResult<()> {
self.exec_plan_with_eval(self::eval)
}
pub fn run_plan_with_eval<Eval, E>(
&mut self,
inputs: TVec<TValue>,
eval: Eval,
) -> TractResult<TVec<TValue>>
where
Eval: for<'a, 'b, 'c> FnMut(
&'a EvalContext<'a>,
Option<&'b mut (dyn OpState + 'static)>,
&'c Node<F, O>,
TVec<TValue>,
) -> Result<TVec<TValue>, E>,
E: Into<anyhow::Error> + Send + Sync + 'static,
{
self.set_inputs(inputs)?;
self.resolve_symbols_with_states()?;
self.exec_plan_with_eval(eval)?;
let outputs = self.outputs()?;
self.reset_turn()?;
Ok(outputs)
}
pub fn exec_plan_with_eval<Eval, E>(&mut self, eval: Eval) -> TractResult<()>
where
Eval: for<'a, 'b, 'c> FnMut(
&'a EvalContext<'a>,
Option<&'b mut (dyn OpState + 'static)>,
&'c Node<F, O>,
TVec<TValue>,
) -> Result<TVec<TValue>, E>,
E: Into<anyhow::Error> + Send + Sync + 'static,
{
if let Some(executor) = self.plan().executor.as_ref() {
tract_linalg::multithread::multithread_tract_scope(executor.clone(), || {
self.do_exec_plan_with_eval(eval)
})
} else {
self.do_exec_plan_with_eval(eval)
}
}
fn do_exec_plan_with_eval<Eval, E>(&mut self, mut eval: Eval) -> TractResult<()>
where
Eval: for<'a, 'b, 'c> FnMut(
&'a EvalContext<'a>,
Option<&'b mut (dyn OpState + 'static)>,
&'c Node<F, O>,
TVec<TValue>,
) -> Result<TVec<TValue>, E>,
E: Into<anyhow::Error> + Send + Sync + 'static,
{
{
self.ready_turn();
self.plan
.turn_handler
.as_ref()
.map(|it| it.before_plan_eval(&mut self.turn_state))
.transpose()?;
let mut syms_done = !self.plan.has_unresolved_symbols
|| self
.plan
.symbols
.iter()
.all(|s| self.turn_state.resolved_symbols.get(s).is_some());
for (step, n) in self.plan.order.iter().enumerate() {
let node = self.plan.model.node(*n);
trace!("Running step {step}, node {node}");
let mut inputs: TVec<TValue> = tvec![];
for i in &node.inputs {
trace!(" use input {i:?}");
let prec_node = self.plan.model.node(i.node);
let prec = self.turn_state.values[i.node].as_ref().ok_or_else(|| {
format_err!("Computing {}, precursor {} not done:", node, prec_node)
})?;
inputs.push(prec[i.slot].clone())
}
for flush in &self.plan.flush_lists[step] {
trace!(" Ran {} can now flush {}", node, self.plan.model.node(*flush));
self.turn_state.values[*flush] = None;
}
if cfg!(debug_assertions) {
let facts = self.plan.model.node_input_facts(node.id)?;
if facts.len() != inputs.len() {
bail!(
"Evaluating {}: expected {} inputs, got {}",
node,
facts.len(),
inputs.len()
);
}
for (ix, (v, f)) in inputs.iter().zip(facts.iter()).enumerate() {
if !f.matches(v, Some(&self.turn_state.resolved_symbols))? {
bail!(
"Evaluating {}: input {:?}, expected {:?}, got {:?}",
node,
ix,
f,
v
);
}
}
}
if node.inputs.is_empty()
&& let Some(preset) = self.turn_state.values[*n].as_ref()
{
inputs = preset.clone();
}
let ctx = self.turn_state.context(self.session, node.id);
let vs = eval(&ctx, self.op_states[node.id].as_deref_mut(), node, inputs)
.map_err(|e| e.into())?;
if !syms_done && self.plan.has_unresolved_symbols {
for (o, v) in node.outputs.iter().zip(vs.iter()) {
if let Ok(f) = o.fact.to_typed_fact() {
for (dim_abstract, dim_concrete) in f.shape.iter().zip(v.shape()) {
Self::resolve(
&mut self.turn_state,
dim_abstract,
*dim_concrete as i64,
)?;
}
}
}
if self
.plan
.symbols
.iter()
.all(|s| self.turn_state.resolved_symbols.get(s).is_some())
{
syms_done = true;
}
}
if cfg!(debug_assertions) {
let facts = self.plan.model.node_output_facts(node.id)?;
if facts.len() != vs.len() {
bail!(
"Evaluating {}: expected {} outputs, got {}",
node,
facts.len(),
vs.len()
);
}
for (ix, (v, f)) in vs.iter().zip(facts.iter()).enumerate() {
if node.outputs[ix].successors.len() == 0 {
continue;
}
if !f.matches(v, Some(&self.turn_state.resolved_symbols))? {
bail!(
"Evaluating {}: output {:?}, expected {:?}, got {:?}",
node,
ix,
f,
v
);
}
}
}
self.turn_state.values[node.id] = Some(vs);
}
self.plan
.turn_handler
.as_ref()
.map(|it| it.after_plan_eval(&mut self.turn_state))
.transpose()?;
}
Ok(())
}
pub fn set_inputs(&mut self, inputs: TVec<TValue>) -> TractResult<()> {
ensure!(
inputs.len() == self.model().inputs.len(),
"Wrong number of inputs for model. Expected {} got {}",
self.model().inputs.len(),
inputs.len()
);
for (ix, t) in inputs.into_iter().enumerate() {
self.set_input(ix, t)?
}
Ok(())
}
pub(crate) fn set_inputs_drain(&mut self, inputs: &mut TVec<TValue>) -> TractResult<()> {
ensure!(
inputs.len() == self.model().inputs.len(),
"Wrong number of inputs for model. Expected {} got {}",
self.model().inputs.len(),
inputs.len()
);
for (ix, t) in inputs.drain(..).enumerate() {
self.set_input(ix, t)?
}
Ok(())
}
fn resolve(state: &mut TurnState, expression: &TDim, provided: i64) -> TractResult<()> {
if let TDim::Sym(sym) = expression
&& state.resolved_symbols.get(sym).is_none()
{
state.resolved_symbols.set(sym, provided);
if state.scenario.is_none() {
let scope = sym.scope().with_context(|| {
format!(
"Symbol {sym:?} points to an invalid (dead ?) SymbolScope. \
Make sure to create symbols using the model-managed SymbolScope."
)
})?;
state.scenario = scope.guess_scenario(&state.resolved_symbols)?;
}
return Ok(());
}
let expected = expression.eval(&state.resolved_symbols);
if let Some(x) = expected.as_i64()
&& x != provided
{
bail!("Clashing resolution for expression. {expression}={x} != {provided}. ({state:?})")
}
if expected.symbols().len() == 1 {
let sym = expected.symbols().into_iter().next().unwrap();
if let Some(v) = solve_for(&sym, &expected, &provided.to_dim()) {
debug!("Determined symbol {sym}={v}");
state.resolved_symbols.set(&sym, v.to_i64().unwrap());
}
if state.scenario.is_none() {
let scope = sym
.scope()
.with_context(|| format!("Symbol {sym:?} points to an invalid (dead ?) SymbolScope. Make sure to create symbols using the model-managed SymbolScope."))?;
state.scenario = scope.guess_scenario(&state.resolved_symbols)?;
}
}
Ok(())
}
pub fn set_input(&mut self, input: usize, t: TValue) -> TractResult<()> {
let outlet: OutletId = *self
.model()
.input_outlets()?
.get(input)
.with_context(|| format!("Invalid input id for model ({input})."))?;
if let Ok(fact) = self.plan.model.outlet_fact(outlet)?.to_typed_fact() {
for (expected, provided) in fact.shape.iter().zip(t.shape()) {
Self::resolve(&mut self.turn_state, expected, *provided as i64)?;
}
}
let fact = self.plan.model.outlet_fact(outlet)?;
ensure!(
fact.matches(&t, Some(&self.turn_state.resolved_symbols))
.with_context(|| format!("Setting input {input}"))?,
"Input at index {input} has incorrect dtype or shape (got {t:?}, expected to match fact {fact:?})",
);
self.ready_turn();
self.turn_state.values[outlet.node] = Some(tvec!(t));
Ok(())
}
pub fn output(&self, id: usize) -> TractResult<&TValue> {
let outlet = self.model().output_outlets()?.get(id).with_context(|| {
format!(
"Required output {}, only have {}",
id,
self.model().output_outlets().unwrap().len()
)
})?;
let value: &TValue = self
.turn_state
.values
.get(outlet.node)
.context("node id for output beyond node values array")?
.as_ref()
.context("node is not an output")?
.get(outlet.slot)
.context("slot id too high")?;
Ok(value)
}
pub fn outputs(&mut self) -> TractResult<TVec<TValue>> {
let &mut SimpleState { ref plan, ref mut turn_state, .. } = self;
let mut v = tvec![];
for o in plan.outputs.iter() {
let vs = turn_state.values[o.node].as_mut().ok_or_else(|| {
format_err!("Outputs of {:?} are not computed", plan.model.nodes()[o.node])
})?;
v.push(vs[o.slot].clone())
}
Ok(v)
}
pub fn set_values(&mut self, id: usize, values: TVec<TValue>) -> TractResult<()> {
self.turn_state.values[id] = Some(values);
Ok(())
}
pub fn set_value(&mut self, id: usize, value: TValue) -> TractResult<()> {
self.set_values(id, tvec!(value))
}
pub fn prepare_inputs(&self, node: usize) -> TractResult<TVec<TValue>> {
let SimpleState { plan, turn_state, .. } = self;
let nodes = plan.model.nodes();
let node = &nodes[node];
let mut inputs: TVec<TValue> = tvec![];
for i in &node.inputs {
let prec_node = &nodes[i.node];
let prec = turn_state.values[i.node].as_ref().ok_or_else(|| {
format_err!("Computing {}, precursor {} not done.", node, prec_node)
})?;
inputs.push(prec[i.slot].clone())
}
if node.inputs.is_empty()
&& let Some(preset) = turn_state.values[node.id].as_ref()
{
inputs = preset.clone();
}
Ok(inputs)
}
pub fn compute_one(&mut self, node: usize) -> TractResult<()> {
let inputs = self.prepare_inputs(node)?;
self.compute_one_with_inputs(node, inputs)
}
pub fn compute_one_with_inputs(
&mut self,
node: usize,
inputs: TVec<TValue>,
) -> TractResult<()> {
let &mut SimpleState {
ref plan,
ref mut turn_state,
op_states: ref mut states,
session,
..
} = self;
let nodes = plan.model.nodes();
let node = &nodes[node];
let ctx = turn_state.context(session, node.id);
let vs = eval(&ctx, states[node.id].as_deref_mut(), node, inputs)?;
turn_state.values[node.id] = Some(vs);
Ok(())
}
pub fn compute_recursively(&mut self, node: usize) -> TractResult<&[TValue]> {
let values = {
#[allow(clippy::needless_collect)] let precs: Vec<usize> =
self.model().nodes()[node].inputs.iter().map(|i| i.node).collect();
for i in precs.into_iter() {
if self.turn_state.values[i].is_none() {
let _ = self.compute_recursively(i)?;
}
}
let mut inputs: TVec<TValue> = tvec![];
{
let node = &self.model().nodes()[node];
for i in &node.inputs {
inputs.push(self.turn_state.values[i.node].as_ref().unwrap()[i.slot].clone())
}
if node.inputs.is_empty()
&& let Some(preset) = self.turn_state.values[node.id].as_ref()
{
inputs = preset.clone();
}
}
let &mut Self {
op_states: ref mut states,
turn_state: ref mut turn,
ref plan,
session,
..
} = self;
let ctx = turn.context(session, node);
eval(&ctx, states[node].as_deref_mut(), &plan.model().nodes[node], inputs)?
};
self.turn_state.values[node] = Some(values);
Ok(self.turn_state.values[node].as_ref().unwrap())
}
pub fn take_by_name(&mut self, name: &str) -> TractResult<TVec<Tensor>> {
let id = self.model().node_by_name(name)?.id;
Self::take(self, id)
}
pub fn take(&mut self, id: usize) -> TractResult<TVec<Tensor>> {
Ok(self.turn_state.values[id]
.take()
.ok_or_else(|| format_err!("Node is not computed"))?
.into_iter()
.map(|v| v.into_tensor())
.collect())
}
pub fn plan(&self) -> &Arc<SimplePlan<F, O>> {
&self.plan
}
pub fn model(&self) -> &Graph<F, O> {
&self.plan.model
}
}
pub fn eval<F, O>(
ctx: &EvalContext,
mut state: Option<&mut (dyn OpState + 'static)>,
node: &Node<F, O>,
input: TVec<TValue>,
) -> TractResult<TVec<TValue>>
where
F: Fact + Clone + 'static,
O: Debug + Display + AsRef<dyn Op> + AsMut<dyn Op> + Clone + 'static,
{
match state {
Some(ref mut state) => state.eval(ctx, node.op(), input),
None => node.op().eval(ctx, input),
}
.with_context(|| format!("Evaluating {node}"))
}
#[cfg(test)]
mod test {
use super::*;
fn is_send<T: Send>() {}
fn is_sync<T: Sync>() {}
#[test]
fn type_model_is_sync() {
is_sync::<TypedModel>();
}
#[test]
fn type_model_is_send() {
is_send::<TypedModel>();
}
#[test]
fn type_plan_is_send() {
is_send::<TypedSimplePlan>();
}
#[test]
fn type_plan_is_sync() {
is_sync::<TypedSimplePlan>();
}
#[test]
fn type_state_is_send() {
is_send::<TypedSimpleState>();
}
#[test]
fn symbol_consumers_run_after_their_definer() -> TractResult<()> {
use crate::ops::array::MultiBroadcastTo;
use crate::ops::array::Range;
let mut model = TypedModel::default();
let r = model.symbols.sym("r").to_dim();
let limit = model.add_source("limit", i64::datum_type().scalar_fact())?;
let start = model.add_const("start", tensor0(0i64))?;
let step = model.add_const("step", tensor0(1i64))?;
let range = model.wire_node("range", Range::new(r.clone()), &[start, limit, step])?[0];
let zero = model.add_const("zero", tensor0(0i64))?;
let expand = model.wire_node(
"expand",
MultiBroadcastTo::new(ShapeFact::from_dims(tvec![r, 2usize.to_dim()])),
&[zero],
)?[0];
model.select_output_outlets(&[expand, range])?;
let found = TypedSimplePlan::new(model)?.run(tvec!(tensor0(5i64).into_tvalue()))?;
assert_eq!(*found[0], tensor2(&[[0i64, 0], [0, 0], [0, 0], [0, 0], [0, 0]]));
assert_eq!(*found[1], tensor1(&[0i64, 1, 2, 3, 4]));
Ok(())
}
#[test]
fn symbol_dep_skipped_when_consumer_is_ancestor() -> TractResult<()> {
use crate::ops::array::MultiBroadcastTo;
use crate::ops::array::Range;
let mut model = TypedModel::default();
let r_a = model.symbols.sym("ra").to_dim();
let r_b = model.symbols.sym("rb").to_dim();
let limit = model.add_source("limit_in", i64::datum_type().scalar_fact())?;
let start = model.add_const("start", tensor0(0i64))?;
let step = model.add_const("step", tensor0(1i64))?;
let range_a =
model.wire_node("range_a", Range::new(r_a.clone()), &[start, limit, step])?[0];
let zero = model.add_const("zero", tensor0(0i64))?;
let expand = model.wire_node(
"expand",
MultiBroadcastTo::new(ShapeFact::from_dims(tvec![r_a.clone(), r_b.clone()])),
&[zero],
)?[0];
let summed = model.wire_node(
"summed",
crate::ops::nn::Reduce::new(tvec!(0, 1), crate::ops::nn::Reducer::Sum),
&[expand],
)?[0];
let summed =
model.wire_node("summed_rm0", crate::ops::change_axes::AxisOp::Rm(0), &[summed])?[0];
let summed =
model.wire_node("summed_rm1", crate::ops::change_axes::AxisOp::Rm(0), &[summed])?[0];
let range_b = model.wire_node("range_b", Range::new(r_b), &[start, summed, step])?;
model.select_output_outlets(&[range_b[0], range_a])?;
TypedSimplePlan::new(model)?;
Ok(())
}
#[test]
fn input_defined_symbol_needs_no_dep() -> TractResult<()> {
use crate::ops::array::MultiBroadcastTo;
let mut model = TypedModel::default();
let s = model.symbols.sym("s").to_dim();
let len = model.add_source("len", i64::datum_type().fact([s.clone()]))?;
let zero = model.add_const("zero", tensor0(0i64))?;
let expand = model.wire_node(
"expand",
MultiBroadcastTo::new(ShapeFact::from_dims(tvec![s, 2usize.to_dim()])),
&[zero],
)?[0];
model.select_output_outlets(&[expand, len])?;
let found = TypedSimplePlan::new(model)?.run(tvec!(tensor1(&[7i64; 7]).into_tvalue()))?;
assert_eq!(found[0].shape(), &[7, 2]);
Ok(())
}
#[test]
fn symbolic_konst_values_consumer_ordered_after_definer() -> TractResult<()> {
use crate::ops::array::Range;
use crate::ops::cast::cast;
use crate::runtime::RunOptions;
let mut model = TypedModel::default();
let r = model.symbols.sym("rk").to_dim();
let limit = model.add_source("limit_in", i64::datum_type().scalar_fact())?;
let start = model.add_const("start", tensor0(0i64))?;
let step = model.add_const("step", tensor0(1i64))?;
let range = model.wire_node("range", Range::new(r.clone()), &[start, limit, step])?[0];
let shape_konst = model.add_const("shape_konst", tensor1(&[r]))?;
let casted = model.wire_node("cast_konst", cast(i64::datum_type()), &[shape_konst])?;
model.select_output_outlets(&[casted[0], range])?;
let options = RunOptions { skip_order_opt_ram: true, ..Default::default() };
let found = TypedSimplePlan::new_with_options(model, &options)?
.run(tvec!(tensor0(4i64).into_tvalue()))?;
assert_eq!(*found[0], tensor1(&[4i64]));
assert_eq!(*found[1], tensor1(&[0i64, 1, 2, 3]));
Ok(())
}
}
#[cfg(test)]
mod symbol_deps_tests {
use super::*;
use crate::ops::array::MultiBroadcastTo;
use crate::ops::array::Range;
fn symbol_deps_of(model: &TypedModel) -> Vec<(usize, usize)> {
let inputs = model.input_outlets().unwrap().iter().map(|n| n.node).collect::<Vec<_>>();
let outputs = model.output_outlets().unwrap().iter().map(|n| n.node).collect::<Vec<_>>();
SimplePlan::<TypedFact, Box<dyn TypedOp>>::symbol_deps(
model.nodes(),
&inputs,
&outputs,
&[],
)
.unwrap()
}
fn range_over(model: &mut TypedModel, sym: &str) -> TractResult<TVec<OutletId>> {
let r = model.symbols.sym(sym).to_dim();
let start = model.add_const("start", tensor0(0i64))?;
let step = model.add_const("step", tensor0(1i64))?;
let limit = model.add_source("limit", i64::datum_type().scalar_fact())?;
model.wire_node("range", Range::new(r), &[start, limit, step])
}
#[test]
fn no_edge_consumer_gets_a_dep() -> TractResult<()> {
let mut model = TypedModel::default();
let range = range_over(&mut model, "r")?[0];
let zero = model.add_const("zero", tensor0(0i64))?;
let r = model.symbols.sym("r").to_dim();
let expand = model.wire_node(
"expand",
MultiBroadcastTo::new(ShapeFact::from_dims(tvec![r, 2usize.to_dim()])),
&[zero],
)?[0];
model.select_output_outlets(&[expand, range])?;
assert_eq!(symbol_deps_of(&model), vec![(expand.node, range.node)]);
Ok(())
}
#[test]
fn input_definers_yield_no_dep() -> TractResult<()> {
let mut model = TypedModel::default();
let s = model.symbols.sym("s").to_dim();
let len = model.add_source("len", i64::datum_type().fact([s.clone()]))?;
let zero = model.add_const("zero", tensor0(0i64))?;
let expand = model.wire_node(
"expand",
MultiBroadcastTo::new(ShapeFact::from_dims(tvec![s, 2usize.to_dim()])),
&[zero],
)?[0];
model.select_output_outlets(&[expand, len])?;
assert!(symbol_deps_of(&model).is_empty());
Ok(())
}
#[test]
fn tensor_ancestor_consumer_gets_no_dep() -> TractResult<()> {
let mut model = TypedModel::default();
let range = range_over(&mut model, "r")?[0];
let casted =
model.wire_node("downstream", crate::ops::cast::cast(f32::datum_type()), &[range])?;
model.select_output_outlets(&[casted[0]])?;
assert!(symbol_deps_of(&model).is_empty());
Ok(())
}
#[test]
fn scan_forwards_body_requirements() -> TractResult<()> {
use crate::ops::array::Range;
use crate::ops::cast::cast;
use crate::ops::konst::Const;
use crate::ops::scan::{InputMapping, OutputMapping, Scan, ScanInfo};
let mut model = TypedModel::default();
let r = model.symbols.sym("b");
let limit = model.add_source("limit", i64::datum_type().scalar_fact())?;
let limit = model.wire_node("limit_tdim", cast(TDim::datum_type()), &[limit])?;
let start =
model.wire_node("start", Const::new(tensor0(TDim::Val(0)).into_arc_tensor())?, &[])?;
let step =
model.wire_node("step", Const::new(tensor0(TDim::Val(1)).into_arc_tensor())?, &[])?;
let range =
model.wire_node("range", Range::new(r.to_dim()), &[start[0], limit[0], step[0]])?;
let mut body = TypedModel { symbols: model.symbols.clone(), ..TypedModel::default() };
let state = body.add_source("state_in", i64::datum_type().fact([1usize]))?;
let scan_in = body.add_source("scan_in", i64::datum_type().fact([1usize]))?;
let konst = body.wire_node(
"body_konst",
Const::new(tensor0(TDim::Sym(r)).into_arc_tensor())?,
&[],
)?;
let _casted = body.wire_node("body_cast", cast(f32::datum_type()), &[konst[0]])?;
let sum = body.wire_node("sum", crate::ops::math::add(), &[state, scan_in])?;
body.select_output_outlets(&[sum[0]])?;
let init = model.add_const("init", tensor1(&[0i64]))?;
let scanned = model.add_source("scanned", i64::datum_type().fact([3usize]))?;
let scan = model.wire_node(
"scan",
Scan::new(
body,
vec![InputMapping::Scan(ScanInfo { axis: 0, chunk: 1 }), InputMapping::State],
vec![OutputMapping {
scan: Some((0, ScanInfo { axis: 0, chunk: 1 })),
full_dim_hint: None,
last_value_slot: None,
state: true,
}],
0,
)?,
&[scanned, init],
)?;
model.select_output_outlets(&[scan[0], range[0]])?;
let deps = symbol_deps_of(&model);
assert!(
deps.contains(&(scan[0].node, range[0].node)),
"expected the scan ordered after the top-level binder, got {deps:?}"
);
Ok(())
}
#[test]
fn hidden_input_konst_mention_yields_dep() -> TractResult<()> {
use crate::ops::array::Range;
use crate::ops::cast::cast;
let mut model = TypedModel::default();
let limit = model.add_source("limit", i64::datum_type().scalar_fact())?;
let limit = model.wire_node("limit_tdim", cast(TDim::datum_type()), &[limit])?;
let start = model.add_const("start", tensor0(TDim::Val(0)))?;
let step = model.add_const("step", tensor0(TDim::Val(1)))?;
let r = model.symbols.sym("h");
let range = model.wire_node("range", Range::new(r.to_dim()), &[start, limit[0], step])?;
let konst = model.wire_node(
"konst",
crate::ops::konst::Const::new(
tensor1(&[TDim::Sym(r.clone()), TDim::Sym(r)]).into_arc_tensor(),
)?,
&[],
)?;
let casted = model.wire_node("cast", cast(f32::datum_type()), &[konst[0]])?;
model.select_output_outlets(&[casted[0], range[0]])?;
let deps = symbol_deps_of(&model);
assert!(
deps.contains(&(casted[0].node, range[0].node)),
"expected the cast ordered after the binder, got {deps:?}"
);
Ok(())
}
#[test]
fn dead_nodes_contribute_nothing() -> TractResult<()> {
let mut model = TypedModel::default();
let range = range_over(&mut model, "r")?[0];
let zero = model.add_const("zero", tensor0(0i64))?;
let r = model.symbols.sym("r").to_dim();
let expand = model.wire_node(
"expand",
MultiBroadcastTo::new(ShapeFact::from_dims(tvec![r.clone(), 2usize.to_dim()])),
&[zero],
)?[0];
model.select_output_outlets(&[expand, range])?;
let before = symbol_deps_of(&model);
model.wire_node(
"dead",
MultiBroadcastTo::new(ShapeFact::from_dims(tvec![r, 3usize.to_dim()])),
&[zero],
)?;
assert_eq!(symbol_deps_of(&model), before);
Ok(())
}
}
#[cfg(test)]
mod late_input_definer_tests {
use super::*;
use crate::ops::array::Range;
#[test]
fn late_input_definer_wins() -> TractResult<()> {
let mut model = TypedModel::default();
let s = model.symbols.sym("S");
let definer = model.add_source("definer", f32::datum_type().fact([s.to_dim()]))?;
let start = model.add_const("start", tensor0(TDim::Val(0)))?;
let step = model.add_const("step", tensor0(TDim::Val(1)))?;
let end = model.add_const("end", tensor0(TDim::Sym(s.clone())))?;
let range = model.wire_node("range", Range::new(s.to_dim()), &[start, end, step])?;
model.select_output_outlets(&[range[0], definer])?;
crate::internal::TypedSimplePlan::new(model)?;
Ok(())
}
}